From 769e12840d53687544e83185714b41a3dc4e8a6e Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Sat, 27 Jun 2026 00:43:07 +0200 Subject: [PATCH] [agent-network] Management: store, manager, synthesizer, policy engine, provider catalog, HTTP/gRPC API Adds the account-scoped agent-network module: provider/policy/budget CRUD and store, the reverse-proxy service synthesizer, policy selection + limit enforcement, the provider catalog (incl. Vertex AI and AWS Bedrock entries), and the management HTTP + proxy gRPC surfaces. --- .../network_map/controller/repository.go | 10 + management/internals/modules/peers/manager.go | 56 +- .../reverseproxy/accesslogs/accesslogentry.go | 5 + .../accesslogs/manager/manager.go | 211 +++- .../modules/reverseproxy/service/service.go | 109 +- management/internals/server/boot.go | 28 +- management/internals/server/modules.go | 19 + management/internals/shared/grpc/proxy.go | 200 ++- management/server/activity/codes.go | 49 + .../server/agentnetwork/catalog/catalog.go | 749 +++++++++++ .../server/agentnetwork/labelgen/labelgen.go | 66 + .../agentnetwork/labelgen/labelgen_test.go | 101 ++ .../server/agentnetwork/labelgen/words.go | 136 ++ management/server/agentnetwork/manager.go | 896 ++++++++++++++ .../server/agentnetwork/policyselect.go | 660 ++++++++++ .../policyselect_account_realstore_test.go | 181 +++ .../policyselect_realstore_test.go | 214 ++++ .../server/agentnetwork/policyselect_test.go | 641 ++++++++++ management/server/agentnetwork/reconcile.go | 131 ++ .../server/agentnetwork/reconcile_test.go | 232 ++++ management/server/agentnetwork/synthesizer.go | 1059 ++++++++++++++++ .../synthesizer_guardrail_realstore_test.go | 178 +++ ...nthesizer_log_collection_realstore_test.go | 70 ++ ...ynthesizer_parser_redact_realstore_test.go | 145 +++ .../synthesizer_realstore_test.go | 174 +++ .../server/agentnetwork/synthesizer_test.go | 1098 +++++++++++++++++ .../server/agentnetwork/types/accesslog.go | 117 ++ .../agentnetwork/types/accesslogfilter.go | 213 ++++ .../server/agentnetwork/types/budgetrule.go | 106 ++ .../server/agentnetwork/types/consumption.go | 69 ++ .../agentnetwork/types/consumption_test.go | 141 +++ .../server/agentnetwork/types/guardrail.go | 120 ++ .../server/agentnetwork/types/policy.go | 192 +++ .../server/agentnetwork/types/provider.go | 252 ++++ .../server/agentnetwork/types/settings.go | 78 ++ management/server/agentnetwork/types/usage.go | 47 + .../agentnetwork/types/usageoverview.go | 96 ++ .../server/agentnetwork/wire_shape_test.go | 109 ++ .../agentnetwork_budgetrule_realstack_test.go | 126 ++ .../agentnetwork_proxypeer_restart_test.go | 199 +++ .../server/agentnetwork_realstack_test.go | 212 ++++ management/server/http/handler.go | 7 +- .../agentnetwork/access_log_handler.go | 91 ++ .../handlers/agentnetwork/budget_handler.go | 172 +++ .../agentnetwork/budget_handler_test.go | 131 ++ .../agentnetwork/consumption_handler.go | 53 + .../agentnetwork/guardrails_handler.go | 171 +++ .../handlers/agentnetwork/handlers_test.go | 256 ++++ .../handlers/agentnetwork/policies_handler.go | 228 ++++ .../agentnetwork/providers_handler.go | 217 ++++ .../handlers/agentnetwork/settings_handler.go | 74 ++ .../testing/testing_tools/channel/channel.go | 4 +- .../server/permissions/modules/module.go | 2 + management/server/store/sql_store.go | 254 ++++ .../server/store/sql_store_agentnetwork.go | 623 ++++++++++ .../sql_store_agentnetwork_accesslog_test.go | 202 +++ .../sql_store_agentnetwork_budgetrule_test.go | 112 ++ management/server/store/store.go | 34 + .../server/store/store_mock_agentnetwork.go | 464 +++++++ 59 files changed, 12576 insertions(+), 14 deletions(-) create mode 100644 management/server/agentnetwork/catalog/catalog.go create mode 100644 management/server/agentnetwork/labelgen/labelgen.go create mode 100644 management/server/agentnetwork/labelgen/labelgen_test.go create mode 100644 management/server/agentnetwork/labelgen/words.go create mode 100644 management/server/agentnetwork/manager.go create mode 100644 management/server/agentnetwork/policyselect.go create mode 100644 management/server/agentnetwork/policyselect_account_realstore_test.go create mode 100644 management/server/agentnetwork/policyselect_realstore_test.go create mode 100644 management/server/agentnetwork/policyselect_test.go create mode 100644 management/server/agentnetwork/reconcile.go create mode 100644 management/server/agentnetwork/reconcile_test.go create mode 100644 management/server/agentnetwork/synthesizer.go create mode 100644 management/server/agentnetwork/synthesizer_guardrail_realstore_test.go create mode 100644 management/server/agentnetwork/synthesizer_log_collection_realstore_test.go create mode 100644 management/server/agentnetwork/synthesizer_parser_redact_realstore_test.go create mode 100644 management/server/agentnetwork/synthesizer_realstore_test.go create mode 100644 management/server/agentnetwork/synthesizer_test.go create mode 100644 management/server/agentnetwork/types/accesslog.go create mode 100644 management/server/agentnetwork/types/accesslogfilter.go create mode 100644 management/server/agentnetwork/types/budgetrule.go create mode 100644 management/server/agentnetwork/types/consumption.go create mode 100644 management/server/agentnetwork/types/consumption_test.go create mode 100644 management/server/agentnetwork/types/guardrail.go create mode 100644 management/server/agentnetwork/types/policy.go create mode 100644 management/server/agentnetwork/types/provider.go create mode 100644 management/server/agentnetwork/types/settings.go create mode 100644 management/server/agentnetwork/types/usage.go create mode 100644 management/server/agentnetwork/types/usageoverview.go create mode 100644 management/server/agentnetwork/wire_shape_test.go create mode 100644 management/server/agentnetwork_budgetrule_realstack_test.go create mode 100644 management/server/agentnetwork_proxypeer_restart_test.go create mode 100644 management/server/agentnetwork_realstack_test.go create mode 100644 management/server/http/handlers/agentnetwork/access_log_handler.go create mode 100644 management/server/http/handlers/agentnetwork/budget_handler.go create mode 100644 management/server/http/handlers/agentnetwork/budget_handler_test.go create mode 100644 management/server/http/handlers/agentnetwork/consumption_handler.go create mode 100644 management/server/http/handlers/agentnetwork/guardrails_handler.go create mode 100644 management/server/http/handlers/agentnetwork/handlers_test.go create mode 100644 management/server/http/handlers/agentnetwork/policies_handler.go create mode 100644 management/server/http/handlers/agentnetwork/providers_handler.go create mode 100644 management/server/http/handlers/agentnetwork/settings_handler.go create mode 100644 management/server/store/sql_store_agentnetwork.go create mode 100644 management/server/store/sql_store_agentnetwork_accesslog_test.go create mode 100644 management/server/store/sql_store_agentnetwork_budgetrule_test.go create mode 100644 management/server/store/store_mock_agentnetwork.go diff --git a/management/internals/controllers/network_map/controller/repository.go b/management/internals/controllers/network_map/controller/repository.go index caef362cb..1b76e6a0b 100644 --- a/management/internals/controllers/network_map/controller/repository.go +++ b/management/internals/controllers/network_map/controller/repository.go @@ -3,7 +3,9 @@ package controller import ( "context" + "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" "github.com/netbirdio/netbird/management/internals/modules/zones" + "github.com/netbirdio/netbird/management/server/agentnetwork" "github.com/netbirdio/netbird/management/server/peer" "github.com/netbirdio/netbird/management/server/store" "github.com/netbirdio/netbird/management/server/types" @@ -16,6 +18,10 @@ type Repository interface { GetPeersByIDs(ctx context.Context, accountID string, peerIDs []string) (map[string]*peer.Peer, error) GetPeerByID(ctx context.Context, accountID string, peerID string) (*peer.Peer, error) GetAccountZones(ctx context.Context, accountID string) ([]*zones.Zone, error) + // SynthesizeAgentNetworkServices returns the in-memory reverse-proxy + // services synthesised from the account's agent-network provider/policy + // state. Empty for accounts without agent-network providers. + SynthesizeAgentNetworkServices(ctx context.Context, accountID string) ([]*service.Service, error) } type repository struct { @@ -50,6 +56,10 @@ func (r *repository) GetPeerByID(ctx context.Context, accountID string, peerID s return r.store.GetPeerByID(ctx, store.LockingStrengthNone, accountID, peerID) } +func (r *repository) SynthesizeAgentNetworkServices(ctx context.Context, accountID string) ([]*service.Service, error) { + return agentnetwork.SynthesizeServices(ctx, r.store, accountID) +} + func (r *repository) GetAccountZones(ctx context.Context, accountID string) ([]*zones.Zone, error) { return r.store.GetAccountZones(ctx, store.LockingStrengthNone, accountID) } diff --git a/management/internals/modules/peers/manager.go b/management/internals/modules/peers/manager.go index e22d1e6e0..239d6b09c 100644 --- a/management/internals/modules/peers/manager.go +++ b/management/internals/modules/peers/manager.go @@ -220,12 +220,36 @@ func (m *managerImpl) GetPeerID(ctx context.Context, peerKey string) (string, er func (m *managerImpl) CreateProxyPeer(ctx context.Context, accountID string, peerKey string, cluster string) error { existingPeerID, err := m.store.GetPeerIDByKey(ctx, store.LockingStrengthNone, peerKey) if err == nil && existingPeerID != "" { - // Peer already exists + // Same pubkey already registered — idempotent. return nil } + // Dedupe stale embedded peer records for the same (account, cluster). + // The proxy generates a fresh WireGuard keypair on every startup + // (proxy/internal/roundtrip/netbird.go), so without this sweep the + // prior embedded peer would linger forever — holding its CGNAT IP + // allocation, polluting other peers' rosters, and (most visibly) + // leaving the synth DNS pointing at the dead address. The + // (account, cluster) tuple identifies "the embedded peer for this + // proxy instance at this cluster"; any record matching that tuple + // with a different pubkey is by definition stale and must go. + staleIDs, err := m.findStaleEmbeddedProxyPeers(ctx, accountID, cluster, peerKey) + if err != nil { + return fmt.Errorf("scan for stale embedded proxy peers: %w", err) + } + if len(staleIDs) > 0 { + // userID="" + checkConnected=false: the deletion is initiated + // by management itself on behalf of the freshly-registering + // proxy, not by an end user; the stale peer may still be + // marked Connected from its prior session, but its session is + // dead by definition (its key no longer exists). + if err := m.DeletePeers(ctx, accountID, staleIDs, "", false); err != nil { + return fmt.Errorf("delete stale embedded proxy peers %v: %w", staleIDs, err) + } + } + name := fmt.Sprintf("proxy-%s", xid.New().String()) - peer := &peer.Peer{ + newPeer := &peer.Peer{ Ephemeral: true, ProxyMeta: peer.ProxyMeta{ Cluster: cluster, @@ -242,10 +266,36 @@ func (m *managerImpl) CreateProxyPeer(ctx context.Context, accountID string, pee }, } - _, _, _, _, err = m.accountManager.AddPeer(ctx, accountID, "", "", peer, true) + _, _, _, _, err = m.accountManager.AddPeer(ctx, accountID, "", "", newPeer, true) if err != nil { return fmt.Errorf("failed to create proxy peer: %w", err) } return nil } + +// findStaleEmbeddedProxyPeers returns the peer IDs of embedded proxy peer +// records in accountID that target the same cluster but carry a different +// WireGuard pubkey than the freshly-registering one. Used by CreateProxyPeer +// to garbage-collect stale records left behind when the proxy restarts with a +// regenerated keypair. +func (m *managerImpl) findStaleEmbeddedProxyPeers(ctx context.Context, accountID, cluster, newKey string) ([]string, error) { + account, err := m.store.GetAccount(ctx, accountID) + if err != nil { + return nil, err + } + var stale []string + for _, p := range account.Peers { + if p == nil || !p.ProxyMeta.Embedded { + continue + } + if p.ProxyMeta.Cluster != cluster { + continue + } + if p.Key == newKey { + continue + } + stale = append(stale, p.ID) + } + return stale, nil +} diff --git a/management/internals/modules/reverseproxy/accesslogs/accesslogentry.go b/management/internals/modules/reverseproxy/accesslogs/accesslogentry.go index f2ecfd5f9..6705eab1a 100644 --- a/management/internals/modules/reverseproxy/accesslogs/accesslogentry.go +++ b/management/internals/modules/reverseproxy/accesslogs/accesslogentry.go @@ -39,6 +39,10 @@ type AccessLogEntry struct { BytesDownload int64 `gorm:"index"` Protocol AccessLogProtocol `gorm:"index"` Metadata map[string]string `gorm:"serializer:json"` + // AgentNetwork marks the entry as emitted by a synthesised agent-network + // service. Sourced from proto.AccessLog.AgentNetwork the proxy stamps + // before shipping. Indexed so the agent-network log surface filters cheaply. + AgentNetwork bool `gorm:"index"` } // FromProto creates an AccessLogEntry from a proto.AccessLog @@ -58,6 +62,7 @@ func (a *AccessLogEntry) FromProto(serviceLog *proto.AccessLog) { a.BytesDownload = serviceLog.GetBytesDownload() a.Protocol = AccessLogProtocol(serviceLog.GetProtocol()) a.Metadata = maps.Clone(serviceLog.GetMetadata()) + a.AgentNetwork = serviceLog.GetAgentNetwork() if sourceIP := serviceLog.GetSourceIp(); sourceIP != "" { if addr, err := netip.ParseAddr(sourceIP); err == nil { diff --git a/management/internals/modules/reverseproxy/accesslogs/manager/manager.go b/management/internals/modules/reverseproxy/accesslogs/manager/manager.go index ced2ec4d1..5f8d80ccd 100644 --- a/management/internals/modules/reverseproxy/accesslogs/manager/manager.go +++ b/management/internals/modules/reverseproxy/accesslogs/manager/manager.go @@ -2,12 +2,15 @@ package manager import ( "context" + "math" + "strconv" "strings" "time" log "github.com/sirupsen/logrus" "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs" + agentNetworkTypes "github.com/netbirdio/netbird/management/server/agentnetwork/types" "github.com/netbirdio/netbird/management/server/geolocation" "github.com/netbirdio/netbird/management/server/permissions" "github.com/netbirdio/netbird/management/server/permissions/modules" @@ -16,6 +19,28 @@ import ( "github.com/netbirdio/netbird/shared/management/status" ) +// Metadata keys the proxy stamps on agent-network access-log entries. These +// mirror the constants in proxy/internal/middleware/keys.go and form the wire +// contract between the proxy and management; management flattens them into +// queryable columns. Keep in sync with the proxy side. +const ( + metaKeyProvider = "llm.provider" + metaKeyModel = "llm.model" + metaKeyResolvedProviderID = "llm.resolved_provider_id" + metaKeySelectedPolicyID = "llm.selected_policy_id" + metaKeyPolicyDecision = "llm_policy.decision" + metaKeyPolicyReason = "llm_policy.reason" + metaKeyInputTokens = "llm.input_tokens" //nolint:gosec // metadata key name, not a credential + metaKeyOutputTokens = "llm.output_tokens" //nolint:gosec // metadata key name, not a credential + metaKeyTotalTokens = "llm.total_tokens" //nolint:gosec // metadata key name, not a credential + metaKeyCostUSDTotal = "cost.usd_total" + metaKeyStream = "llm.stream" + metaKeySessionID = "llm.session_id" + metaKeyAuthorisingGroups = "llm.authorising_groups" + metaKeyRequestPrompt = "llm.request_prompt" + metaKeyResponseCompletion = "llm.response_completion" +) + type managerImpl struct { store store.Store permissionsManager permissions.Manager @@ -31,8 +56,14 @@ func NewManager(store store.Store, permissionsManager permissions.Manager, geo g } } -// SaveAccessLog saves an access log entry to the database after enriching it +// SaveAccessLog saves an access log entry to the database after enriching it. +// Agent-network entries are flattened into their own dedicated table (queryable +// LLM columns + group child rows) instead of the shared reverse-proxy table. func (m *managerImpl) SaveAccessLog(ctx context.Context, logEntry *accesslogs.AccessLogEntry) error { + if logEntry.AgentNetwork { + return m.saveAgentNetworkAccessLog(ctx, logEntry) + } + if m.geo != nil && logEntry.GeoLocation.ConnectionIP != nil { location, err := m.geo.Lookup(logEntry.GeoLocation.ConnectionIP) if err != nil { @@ -61,6 +92,184 @@ func (m *managerImpl) SaveAccessLog(ctx context.Context, logEntry *accesslogs.Ac return nil } +// saveAgentNetworkAccessLog flattens the metadata-bearing access-log entry and +// persists it in two parts: +// +// - The stripped usage record is written unconditionally — usage/cost is +// collected on every request regardless of the account's log-collection +// toggle (the proxy ships a usage-only entry when logging is disabled). +// - The full access-log row (with request detail + prompt) is written only +// when the account's EnableLogCollection setting is on. This setting read +// is the authoritative gate; the proxy-side strip is defense in depth. +func (m *managerImpl) saveAgentNetworkAccessLog(ctx context.Context, logEntry *accesslogs.AccessLogEntry) error { + entry, groups := flattenAgentNetworkLog(logEntry) + + usage, usageGroups := usageFromFlattenedLog(entry, groups) + if err := m.store.CreateAgentNetworkUsage(ctx, usage, usageGroups); err != nil { + log.WithContext(ctx).WithFields(log.Fields{ + "account_id": entry.AccountID, + "model": entry.Model, + }).Errorf("failed to save agent-network usage: %v", err) + return err + } + + settings, err := m.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, entry.AccountID) + if err != nil { + // No settings row (or a transient read error) means we can't confirm + // log collection is enabled — usage is already saved, so skip the full + // row rather than fail the whole ingest. + log.WithContext(ctx).Debugf("skipping full agent-network access-log row for account %s: %v", entry.AccountID, err) + return nil + } + if !settings.EnableLogCollection { + return nil + } + + if err := m.store.CreateAgentNetworkAccessLog(ctx, entry, groups); err != nil { + log.WithContext(ctx).WithFields(log.Fields{ + "account_id": entry.AccountID, + "service_id": entry.ServiceID, + "model": entry.Model, + "status": entry.StatusCode, + }).Errorf("failed to save agent-network access log: %v", err) + return err + } + return nil +} + +// flattenAgentNetworkLog converts a reverse-proxy AccessLogEntry (whose LLM +// dimensions live in the opaque Metadata map) into the flattened +// agent-network row + authorising-group child rows. +func flattenAgentNetworkLog(e *accesslogs.AccessLogEntry) (*agentNetworkTypes.AgentNetworkAccessLog, []agentNetworkTypes.AgentNetworkAccessLogGroup) { + meta := e.Metadata + + var sourceIP string + if e.GeoLocation.ConnectionIP != nil { + sourceIP = e.GeoLocation.ConnectionIP.String() + } + + entry := &agentNetworkTypes.AgentNetworkAccessLog{ + ID: e.ID, + AccountID: e.AccountID, + ServiceID: e.ServiceID, + Timestamp: e.Timestamp, + UserID: e.UserId, + SourceIP: sourceIP, + Method: e.Method, + Host: e.Host, + Path: e.Path, + Duration: e.Duration, + StatusCode: e.StatusCode, + AuthMethod: e.AuthMethodUsed, + BytesUpload: e.BytesUpload, + BytesDownload: e.BytesDownload, + + Provider: meta[metaKeyProvider], + Model: meta[metaKeyModel], + SessionID: meta[metaKeySessionID], + ResolvedProviderID: meta[metaKeyResolvedProviderID], + SelectedPolicyID: meta[metaKeySelectedPolicyID], + Decision: meta[metaKeyPolicyDecision], + DenyReason: meta[metaKeyPolicyReason], + InputTokens: parseMetaInt(meta, metaKeyInputTokens), + OutputTokens: parseMetaInt(meta, metaKeyOutputTokens), + TotalTokens: parseMetaInt(meta, metaKeyTotalTokens), + CostUSD: parseMetaFloat(meta, metaKeyCostUSDTotal), + Stream: parseMetaBool(meta, metaKeyStream), + RequestPrompt: meta[metaKeyRequestPrompt], + ResponseCompletion: meta[metaKeyResponseCompletion], + } + + var groups []agentNetworkTypes.AgentNetworkAccessLogGroup + for _, gid := range parseGroupCSV(meta[metaKeyAuthorisingGroups]) { + groups = append(groups, agentNetworkTypes.AgentNetworkAccessLogGroup{ + LogID: entry.ID, + GroupID: gid, + AccountID: entry.AccountID, + }) + } + return entry, groups +} + +// usageFromFlattenedLog derives the stripped usage record (and its group child +// rows) from an already-flattened access-log entry. The usage row shares the +// log's ID so the two correlate. +func usageFromFlattenedLog(e *agentNetworkTypes.AgentNetworkAccessLog, groups []agentNetworkTypes.AgentNetworkAccessLogGroup) (*agentNetworkTypes.AgentNetworkUsage, []agentNetworkTypes.AgentNetworkUsageGroup) { + usage := &agentNetworkTypes.AgentNetworkUsage{ + ID: e.ID, + AccountID: e.AccountID, + Timestamp: e.Timestamp, + UserID: e.UserID, + ResolvedProviderID: e.ResolvedProviderID, + Provider: e.Provider, + Model: e.Model, + SessionID: e.SessionID, + InputTokens: e.InputTokens, + OutputTokens: e.OutputTokens, + TotalTokens: e.TotalTokens, + CostUSD: e.CostUSD, + } + + usageGroups := make([]agentNetworkTypes.AgentNetworkUsageGroup, 0, len(groups)) + for _, g := range groups { + usageGroups = append(usageGroups, agentNetworkTypes.AgentNetworkUsageGroup{ + UsageID: usage.ID, + GroupID: g.GroupID, + AccountID: g.AccountID, + }) + } + return usage, usageGroups +} + +// parseMetaInt parses a non-negative token count. Negative or unparseable +// values are clamped to 0 so a malformed metric can't persist a negative +// counter. +func parseMetaInt(meta map[string]string, key string) int64 { + if v, err := strconv.ParseInt(strings.TrimSpace(meta[key]), 10, 64); err == nil && v >= 0 { + return v + } + return 0 +} + +// parseMetaFloat parses a non-negative, finite cost. Negative, NaN, Inf, or +// unparseable values are clamped to 0 so a malformed metric can't poison the +// stored cost. +func parseMetaFloat(meta map[string]string, key string) float64 { + if v, err := strconv.ParseFloat(strings.TrimSpace(meta[key]), 64); err == nil && v >= 0 && !math.IsInf(v, 0) { + return v + } + return 0 +} + +func parseMetaBool(meta map[string]string, key string) bool { + v, _ := strconv.ParseBool(strings.TrimSpace(meta[key])) + return v +} + +// parseGroupCSV splits the comma-separated authorising-group id list the proxy +// emits, trimming blanks and de-duplicating. Dedup matters because the group +// rows are keyed by (log_id, group_id) / (usage_id, group_id): a repeated id +// in the CSV would otherwise produce a duplicate primary key and fail the +// insert transaction. +func parseGroupCSV(raw string) []string { + if raw == "" { + return nil + } + parts := strings.Split(raw, ",") + out := make([]string, 0, len(parts)) + seen := make(map[string]struct{}, len(parts)) + for _, p := range parts { + if p = strings.TrimSpace(p); p != "" { + if _, dup := seen[p]; dup { + continue + } + seen[p] = struct{}{} + out = append(out, p) + } + } + return out +} + // GetAllAccessLogs retrieves access logs for an account with pagination and filtering func (m *managerImpl) GetAllAccessLogs(ctx context.Context, accountID, userID string, filter *accesslogs.AccessLogFilter) ([]*accesslogs.AccessLogEntry, int64, error) { ok, ctx, err := m.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Services, operations.Read) diff --git a/management/internals/modules/reverseproxy/service/service.go b/management/internals/modules/reverseproxy/service/service.go index ee1e3c8b2..b6438abde 100644 --- a/management/internals/modules/reverseproxy/service/service.go +++ b/management/internals/modules/reverseproxy/service/service.go @@ -66,6 +66,51 @@ type TargetOptions struct { // reachable without WireGuard (public APIs, LAN services, localhost // sidecars). Default false. DirectUpstream bool `json:"direct_upstream,omitempty"` + // Middlewares carries per-target agent-network middleware configs. Empty + // for private and operator-defined services; populated only by the + // agent-network synthesizer. + Middlewares []MiddlewareConfig `gorm:"serializer:json" json:"middlewares,omitempty"` + CaptureMaxRequestBytes int64 `json:"capture_max_request_bytes,omitempty"` + CaptureMaxResponseBytes int64 `json:"capture_max_response_bytes,omitempty"` + CaptureContentTypes []string `gorm:"serializer:json" json:"capture_content_types,omitempty"` + // AgentNetwork marks targets synthesised from Agent Network state. The + // proxy uses it to gate agent-network-specific behaviour (access log + // tagging, observability, etc.). + AgentNetwork bool `json:"agent_network,omitempty"` + // DisableAccessLog suppresses the per-request access-log emission for this + // target. Defaults false to preserve access-log behaviour for every + // non-agent-network target. The agent-network synthesizer sets this true + // only when the account's EnableLogCollection toggle is off. + DisableAccessLog bool `json:"disable_access_log,omitempty"` +} + +// MiddlewareSlot mirrors proto.MiddlewareSlot / middleware.Slot. +type MiddlewareSlot string + +const ( + MiddlewareSlotOnRequest MiddlewareSlot = "on_request" + MiddlewareSlotOnResponse MiddlewareSlot = "on_response" + MiddlewareSlotTerminal MiddlewareSlot = "terminal" +) + +// MiddlewareFailMode mirrors proto.MiddlewareConfig_FailMode. +type MiddlewareFailMode string + +const ( + MiddlewareFailOpen MiddlewareFailMode = "fail_open" + MiddlewareFailClosed MiddlewareFailMode = "fail_closed" +) + +// MiddlewareConfig is the per-target configuration for a single +// middleware instance. Mirrors proto.MiddlewareConfig. +type MiddlewareConfig struct { + ID string `json:"id"` + Enabled bool `json:"enabled"` + Slot MiddlewareSlot `json:"slot"` + ConfigJSON []byte `json:"config_json,omitempty"` + FailMode MiddlewareFailMode `json:"fail_mode,omitempty"` + TimeoutMs int32 `json:"timeout_ms,omitempty"` + CanMutate bool `json:"can_mutate"` } type Target struct { @@ -504,21 +549,75 @@ func targetOptionsToAPI(opts TargetOptions) *api.ServiceTargetOptions { func targetOptionsToProto(opts TargetOptions) *proto.PathTargetOptions { if !opts.SkipTLSVerify && opts.PathRewrite == "" && opts.RequestTimeout == 0 && - len(opts.CustomHeaders) == 0 && !opts.DirectUpstream { + len(opts.CustomHeaders) == 0 && !opts.DirectUpstream && + len(opts.Middlewares) == 0 && opts.CaptureMaxRequestBytes == 0 && + opts.CaptureMaxResponseBytes == 0 && len(opts.CaptureContentTypes) == 0 && + !opts.AgentNetwork && !opts.DisableAccessLog { return nil } popts := &proto.PathTargetOptions{ - SkipTlsVerify: opts.SkipTLSVerify, - PathRewrite: pathRewriteToProto(opts.PathRewrite), - CustomHeaders: opts.CustomHeaders, - DirectUpstream: opts.DirectUpstream, + SkipTlsVerify: opts.SkipTLSVerify, + PathRewrite: pathRewriteToProto(opts.PathRewrite), + CustomHeaders: opts.CustomHeaders, + DirectUpstream: opts.DirectUpstream, + AgentNetwork: opts.AgentNetwork, + DisableAccessLog: opts.DisableAccessLog, } if opts.RequestTimeout != 0 { popts.RequestTimeout = durationpb.New(opts.RequestTimeout) } + if len(opts.Middlewares) > 0 { + popts.Middlewares = middlewaresToProto(opts.Middlewares) + } + popts.CaptureMaxRequestBytes = opts.CaptureMaxRequestBytes + popts.CaptureMaxResponseBytes = opts.CaptureMaxResponseBytes + if len(opts.CaptureContentTypes) > 0 { + popts.CaptureContentTypes = append([]string(nil), opts.CaptureContentTypes...) + } return popts } +// middlewaresToProto converts the internal middleware slice to the proto +// representation sent to the proxy via the mapping stream. +func middlewaresToProto(in []MiddlewareConfig) []*proto.MiddlewareConfig { + out := make([]*proto.MiddlewareConfig, 0, len(in)) + for _, m := range in { + pm := &proto.MiddlewareConfig{ + Id: m.ID, + Enabled: m.Enabled, + Slot: middlewareSlotToProto(m.Slot), + ConfigJson: append([]byte(nil), m.ConfigJSON...), + CanMutate: m.CanMutate, + FailMode: middlewareFailModeToProto(m.FailMode), + } + if m.TimeoutMs > 0 { + pm.Timeout = durationpb.New(time.Duration(m.TimeoutMs) * time.Millisecond) + } + out = append(out, pm) + } + return out +} + +func middlewareSlotToProto(s MiddlewareSlot) proto.MiddlewareSlot { + switch s { + case MiddlewareSlotOnRequest: + return proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST + case MiddlewareSlotOnResponse: + return proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_RESPONSE + case MiddlewareSlotTerminal: + return proto.MiddlewareSlot_MIDDLEWARE_SLOT_TERMINAL + default: + return proto.MiddlewareSlot_MIDDLEWARE_SLOT_UNSPECIFIED + } +} + +func middlewareFailModeToProto(m MiddlewareFailMode) proto.MiddlewareConfig_FailMode { + if m == MiddlewareFailClosed { + return proto.MiddlewareConfig_FAIL_CLOSED + } + return proto.MiddlewareConfig_FAIL_OPEN +} + // l4TargetOptionsToProto converts L4-relevant target options to proto. func l4TargetOptionsToProto(target *Target) *proto.PathTargetOptions { if !target.ProxyProtocol && target.Options.RequestTimeout == 0 && target.Options.SessionIdleTimeout == 0 { diff --git a/management/internals/server/boot.go b/management/internals/server/boot.go index ae82b60fe..82d18b591 100644 --- a/management/internals/server/boot.go +++ b/management/internals/server/boot.go @@ -26,9 +26,11 @@ import ( "github.com/netbirdio/netbird/formatter/hook" "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs" accesslogsmanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs/manager" + rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc" "github.com/netbirdio/netbird/management/server/activity" activitystore "github.com/netbirdio/netbird/management/server/activity/store" + "github.com/netbirdio/netbird/management/server/agentnetwork" nbcache "github.com/netbirdio/netbird/management/server/cache" nbContext "github.com/netbirdio/netbird/management/server/context" nbhttp "github.com/netbirdio/netbird/management/server/http" @@ -120,7 +122,7 @@ func (s *BaseServer) EventStore() activity.Store { func (s *BaseServer) APIHandler() http.Handler { return Create(s, func() http.Handler { - httpAPIHandler, err := nbhttp.NewAPIHandler(context.Background(), s.Router(), s.AccountManager(), s.NetworksManager(), s.ResourcesManager(), s.RoutesManager(), s.GroupsManager(), s.GeoLocationManager(), s.AuthManager(), s.Metrics(), s.PermissionsManager(), s.SettingsManager(), s.ZonesManager(), s.RecordsManager(), s.NetworkMapController(), s.IdpManager(), s.ServiceManager(), s.ReverseProxyDomainManager(), s.AccessLogsManager(), s.ReverseProxyGRPCServer(), s.Config.ReverseProxy.TrustedHTTPProxies, s.RateLimiter(), s.IsValidChildAccount) + httpAPIHandler, err := nbhttp.NewAPIHandler(context.Background(), s.Router(), s.AccountManager(), s.NetworksManager(), s.ResourcesManager(), s.RoutesManager(), s.GroupsManager(), s.GeoLocationManager(), s.AuthManager(), s.Metrics(), s.PermissionsManager(), s.SettingsManager(), s.ZonesManager(), s.RecordsManager(), s.NetworkMapController(), s.IdpManager(), s.ServiceManager(), s.ReverseProxyDomainManager(), s.AccessLogsManager(), s.ReverseProxyGRPCServer(), s.Config.ReverseProxy.TrustedHTTPProxies, s.RateLimiter(), s.IsValidChildAccount, s.AgentNetworkManager()) if err != nil { log.Fatalf("failed to create API handler: %v", err) } @@ -223,11 +225,35 @@ func (s *BaseServer) ReverseProxyGRPCServer() *nbgrpc.ProxyServiceServer { s.AfterInit(func(s *BaseServer) { proxyService.SetServiceManager(s.ServiceManager()) proxyService.SetProxyController(s.ServiceProxyController()) + proxyService.SetAgentNetworkSynthesizer(newAgentNetworkSynthesizer(s.Store())) + proxyService.SetAgentNetworkLimitsService(s.AgentNetworkManager()) }) return proxyService }) } +// agentNetworkSynthesizerAdapter implements nbgrpc.AgentNetworkSynthesizer by +// delegating to the agentnetwork package's store-backed synthesiser. +type agentNetworkSynthesizerAdapter struct { + store store.Store +} + +func newAgentNetworkSynthesizer(s store.Store) *agentNetworkSynthesizerAdapter { + return &agentNetworkSynthesizerAdapter{store: s} +} + +func (a *agentNetworkSynthesizerAdapter) SynthesizeServicesForCluster(ctx context.Context, clusterAddr string) ([]*rpservice.Service, error) { + return agentnetwork.SynthesizeServicesForCluster(ctx, a.store, clusterAddr) +} + +func (a *agentNetworkSynthesizerAdapter) SynthesizeServicesForAccount(ctx context.Context, accountID string) ([]*rpservice.Service, error) { + return agentnetwork.SynthesizeServices(ctx, a.store, accountID) +} + +func (a *agentNetworkSynthesizerAdapter) SynthesizeServiceForDomain(ctx context.Context, domain string) (*rpservice.Service, error) { + return agentnetwork.SynthesizeServiceForDomain(ctx, a.store, domain) +} + func (s *BaseServer) proxyOIDCConfig() nbgrpc.ProxyOIDCConfig { return Create(s, func() nbgrpc.ProxyOIDCConfig { return nbgrpc.ProxyOIDCConfig{ diff --git a/management/internals/server/modules.go b/management/internals/server/modules.go index a70da855a..d99af3971 100644 --- a/management/internals/server/modules.go +++ b/management/internals/server/modules.go @@ -20,6 +20,7 @@ import ( recordsManager "github.com/netbirdio/netbird/management/internals/modules/zones/records/manager" "github.com/netbirdio/netbird/management/server" "github.com/netbirdio/netbird/management/server/account" + "github.com/netbirdio/netbird/management/server/agentnetwork" "github.com/netbirdio/netbird/management/server/geolocation" "github.com/netbirdio/netbird/management/server/groups" "github.com/netbirdio/netbird/management/server/idp" @@ -194,6 +195,24 @@ func (s *BaseServer) NetworksManager() networks.Manager { }) } +func (s *BaseServer) AgentNetworkManager() agentnetwork.Manager { + return Create(s, func() agentnetwork.Manager { + mgr := agentnetwork.NewManager( + s.Store(), + s.PermissionsManager(), + s.AccountManager(), + s.ServiceProxyController(), + ) + // Sweep expired agent-network access logs per account retention, + // reusing the reverse-proxy cleanup interval config. + mgr.StartAccessLogCleanup( + context.Background(), + s.Config.ReverseProxy.AccessLogCleanupIntervalHours, + ) + return mgr + }) +} + func (s *BaseServer) ZonesManager() zones.Manager { return Create(s, func() zones.Manager { return zonesManager.NewManager(s.Store(), s.AccountManager(), s.PermissionsManager(), s.DNSDomain()) diff --git a/management/internals/shared/grpc/proxy.go b/management/internals/shared/grpc/proxy.go index 76663f898..336d74e7e 100644 --- a/management/internals/shared/grpc/proxy.go +++ b/management/internals/shared/grpc/proxy.go @@ -10,6 +10,7 @@ import ( "errors" "fmt" "io" + "math" "net" "net/http" "net/url" @@ -35,6 +36,7 @@ import ( "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey" "github.com/netbirdio/netbird/management/server/idp" "github.com/netbirdio/netbird/management/server/peer" + "github.com/netbirdio/netbird/management/server/agentnetwork" "github.com/netbirdio/netbird/management/server/types" "github.com/netbirdio/netbird/management/server/users" proxyauth "github.com/netbirdio/netbird/proxy/auth" @@ -60,6 +62,23 @@ type ProxyTokenChecker interface { } // ProxyServiceServer implements the ProxyService gRPC server +// AgentNetworkSynthesizer produces in-memory reverse-proxy services from +// Agent Network provider/policy state for the proxy snapshot path; synthesised +// services never appear in the reverseproxy_services table. +type AgentNetworkSynthesizer interface { + SynthesizeServicesForCluster(ctx context.Context, clusterAddr string) ([]*rpservice.Service, error) + SynthesizeServicesForAccount(ctx context.Context, accountID string) ([]*rpservice.Service, error) + SynthesizeServiceForDomain(ctx context.Context, domain string) (*rpservice.Service, error) +} + +// AgentNetworkLimitsService is the minimal slice of agentnetwork.Manager the +// gRPC layer needs for CheckLLMPolicyLimits + RecordLLMUsage — kept narrow so +// the grpc package doesn't take a hard import on the full manager. +type AgentNetworkLimitsService interface { + SelectPolicyForRequest(ctx context.Context, in agentnetwork.PolicySelectionInput) (*agentnetwork.PolicySelectionResult, error) + RecordUsage(ctx context.Context, in agentnetwork.RecordUsageInput) error +} + type ProxyServiceServer struct { proto.UnimplementedProxyServiceServer @@ -72,6 +91,14 @@ type ProxyServiceServer struct { mu sync.RWMutex // Manager for reverse proxy operations serviceManager rpservice.Manager + // agentNetworkSynth produces synthesised reverse-proxy services from + // Agent Network state. Optional — when nil the snapshot path only ships + // persisted services. + agentNetworkSynth AgentNetworkSynthesizer + // agentNetworkLimits handles the pre-flight selection (CheckLLMPolicyLimits) + // and the post-flight consumption write (RecordLLMUsage). Optional — when + // nil both RPCs return Unimplemented. + agentNetworkLimits AgentNetworkLimitsService // ProxyController for service updates and cluster management proxyController proxy.Controller @@ -209,6 +236,127 @@ func (s *ProxyServiceServer) SetServiceManager(manager rpservice.Manager) { s.serviceManager = manager } +// SetAgentNetworkSynthesizer wires the agent-network service synthesiser. +// Optional — when nil the snapshot path skips agent-network synthesis. The +// modules layer injects this after both the proxy server and the agent-network +// manager are constructed. +func (s *ProxyServiceServer) SetAgentNetworkSynthesizer(synth AgentNetworkSynthesizer) { + s.mu.Lock() + s.agentNetworkSynth = synth + s.mu.Unlock() +} + +// SetAgentNetworkLimitsService wires the policy-selection + post-flight +// consumption sink. Pass nil to disable; both RPCs return Unimplemented while +// unset so partial wiring surfaces during integration. +func (s *ProxyServiceServer) SetAgentNetworkLimitsService(svc AgentNetworkLimitsService) { + s.mu.Lock() + s.agentNetworkLimits = svc + s.mu.Unlock() +} + +// agentNetworkSynthesizer returns the synthesiser under read lock. +func (s *ProxyServiceServer) agentNetworkSynthesizer() AgentNetworkSynthesizer { + s.mu.RLock() + defer s.mu.RUnlock() + return s.agentNetworkSynth +} + +// CheckLLMPolicyLimits is the pre-flight policy gate the proxy calls before +// forwarding an LLM request upstream. Delegates to the agent-network selector, +// which scores applicable policies by remaining headroom and returns the +// policy that pays for this request (or a deny when all are exhausted). +func (s *ProxyServiceServer) CheckLLMPolicyLimits(ctx context.Context, req *proto.CheckLLMPolicyLimitsRequest) (*proto.CheckLLMPolicyLimitsResponse, error) { + s.mu.RLock() + svc := s.agentNetworkLimits + s.mu.RUnlock() + if svc == nil { + return nil, status.Errorf(codes.Unimplemented, "agent-network limits service not configured on management") + } + if req.GetAccountId() == "" { + return nil, status.Errorf(codes.InvalidArgument, "account_id is required") + } + if err := enforceAccountScope(ctx, req.GetAccountId()); err != nil { + return nil, err + } + + res, err := svc.SelectPolicyForRequest(ctx, agentnetwork.PolicySelectionInput{ + AccountID: req.GetAccountId(), + UserID: req.GetUserId(), + GroupIDs: req.GetGroupIds(), + ProviderID: req.GetProviderId(), + }) + if err != nil { + log.WithContext(ctx).Errorf("select policy for request: %v", err) + return nil, status.Error(codes.Internal, "select policy failed") + } + + if !res.Allow { + return &proto.CheckLLMPolicyLimitsResponse{ + Decision: "deny", + SelectedPolicyId: res.SelectedPolicyID, + AttributionGroupId: res.AttributionGroupID, + WindowSeconds: res.WindowSeconds, + DenyCode: res.DenyCode, + DenyReason: res.DenyReason, + }, nil + } + return &proto.CheckLLMPolicyLimitsResponse{ + Decision: "allow", + SelectedPolicyId: res.SelectedPolicyID, + AttributionGroupId: res.AttributionGroupID, + WindowSeconds: res.WindowSeconds, + }, nil +} + +// RecordLLMUsage increments the per-(dimension, window) consumption counter for +// the user and optional attribution group after a served request. Returns +// Unimplemented when the agent-network limits service hasn't been wired. +func (s *ProxyServiceServer) RecordLLMUsage(ctx context.Context, req *proto.RecordLLMUsageRequest) (*proto.RecordLLMUsageResponse, error) { + s.mu.RLock() + svc := s.agentNetworkLimits + s.mu.RUnlock() + if svc == nil { + return nil, status.Errorf(codes.Unimplemented, "agent-network limits service not configured on management") + } + + accountID := req.GetAccountId() + if accountID == "" { + return nil, status.Errorf(codes.InvalidArgument, "account_id is required") + } + if err := enforceAccountScope(ctx, accountID); err != nil { + return nil, err + } + tokensIn := req.GetTokensInput() + tokensOut := req.GetTokensOutput() + costUSD := req.GetCostUsd() + + // Reject impossible counters at the boundary instead of recording them: + // a negative window, negative tokens, or a negative / non-finite cost + // would otherwise decrement or poison the persisted consumption totals. + if req.GetWindowSeconds() < 0 || tokensIn < 0 || tokensOut < 0 || costUSD < 0 || math.IsNaN(costUSD) || math.IsInf(costUSD, 0) { + return nil, status.Errorf(codes.InvalidArgument, "usage counters must be non-negative and finite") + } + + // Book the policy-window dimensions (when a policy cap bound this request) + // and every applicable account budget rule's window in a single batched + // transaction. + if err := svc.RecordUsage(ctx, agentnetwork.RecordUsageInput{ + AccountID: accountID, + UserID: req.GetUserId(), + AttributionGroupID: req.GetGroupId(), + GroupIDs: req.GetGroupIds(), + WindowSeconds: req.GetWindowSeconds(), + TokensIn: tokensIn, + TokensOut: tokensOut, + CostUSD: costUSD, + }); err != nil { + log.WithContext(ctx).Errorf("record usage: %v", err) + return nil, status.Error(codes.Internal, "record usage failed") + } + return &proto.RecordLLMUsageResponse{}, nil +} + // SetProxyController sets the proxy controller. Must be called before serving. func (s *ProxyServiceServer) SetProxyController(proxyController proxy.Controller) { s.mu.Lock() @@ -623,12 +771,40 @@ func (s *ProxyServiceServer) snapshotServiceMappings(ctx context.Context, conn * return nil, fmt.Errorf("get services from store: %w", err) } + if synth := s.agentNetworkSynthesizer(); synth != nil { + var synthesised []*rpservice.Service + var serr error + // Account-scoped connections synthesise only their own account, so the + // snapshot can never carry another tenant's mappings (which embed the + // upstream auth header derived from that tenant's provider API key). + // Global connections still see the whole cluster. + if conn.accountID != nil { + synthesised, serr = synth.SynthesizeServicesForAccount(ctx, *conn.accountID) + } else { + synthesised, serr = synth.SynthesizeServicesForCluster(ctx, conn.address) + } + if serr != nil { + // Surface a real synthesis failure instead of silently shipping an + // incomplete snapshot (which would drop the account's agent-network + // routes). Consistent with the persisted-services error above; the + // proxy retries the snapshot on connection error. + return nil, fmt.Errorf("synthesise agent-network services: %w", serr) + } + services = append(services, synthesised...) + } + oidcCfg := s.GetOIDCValidationConfig() var mappings []*proto.ProxyMapping for _, service := range services { if !service.Enabled || service.ProxyCluster == "" || service.ProxyCluster != conn.address { continue } + // Defense in depth: an account-scoped proxy must never receive another + // account's mapping, matching the per-account filtering the incremental + // update path already applies. + if conn.accountID != nil && service.AccountID != *conn.accountID { + continue + } m := service.ToProtoMapping(rpservice.Create, "", oidcCfg) if !proxyAcceptsMapping(conn, m) { @@ -1617,7 +1793,29 @@ func (s *ProxyServiceServer) ValidateSession(ctx context.Context, req *proto.Val } func (s *ProxyServiceServer) getServiceByDomain(ctx context.Context, domain string) (*rpservice.Service, error) { - return s.serviceManager.GetServiceByDomain(ctx, domain) + service, err := s.serviceManager.GetServiceByDomain(ctx, domain) + if err == nil { + return service, nil + } + + // Fall back to the Agent Network synthesiser scoped directly to the domain's + // account. Synthesised services are never persisted, so they must resolve + // here for OIDC / session / tunnel-peer flows against agent-network + // endpoints. Resolving by domain synthesises only the owning account rather + // than every tenant on the cluster. + if synth := s.agentNetworkSynthesizer(); synth != nil { + svc, serr := synth.SynthesizeServiceForDomain(ctx, domain) + if serr != nil { + // A real synthesis failure must surface, not be masked by the + // original store miss — otherwise a transient DB error looks like + // "no such service". + return nil, fmt.Errorf("synthesize agent-network service for %s: %w", domain, serr) + } + if svc != nil { + return svc, nil + } + } + return nil, err } func (s *ProxyServiceServer) checkGroupAccess(service *rpservice.Service, user *types.User) error { diff --git a/management/server/activity/codes.go b/management/server/activity/codes.go index 852193a3b..cfd809871 100644 --- a/management/server/activity/codes.go +++ b/management/server/activity/codes.go @@ -245,6 +245,37 @@ const ( // tunnel. Distinct from UserLoggedInPeer (full interactive login). UserExtendedPeerSession Activity = 125 + // AgentNetworkProviderCreated indicates that a user created an Agent Network provider + AgentNetworkProviderCreated Activity = 126 + // AgentNetworkProviderUpdated indicates that a user updated an Agent Network provider + AgentNetworkProviderUpdated Activity = 127 + // AgentNetworkProviderDeleted indicates that a user deleted an Agent Network provider + AgentNetworkProviderDeleted Activity = 128 + + // AgentNetworkPolicyCreated indicates that a user created an Agent Network policy + AgentNetworkPolicyCreated Activity = 129 + // AgentNetworkPolicyUpdated indicates that a user updated an Agent Network policy + AgentNetworkPolicyUpdated Activity = 130 + // AgentNetworkPolicyDeleted indicates that a user deleted an Agent Network policy + AgentNetworkPolicyDeleted Activity = 131 + + // AgentNetworkGuardrailCreated indicates that a user created an Agent Network guardrail + AgentNetworkGuardrailCreated Activity = 132 + // AgentNetworkGuardrailUpdated indicates that a user updated an Agent Network guardrail + AgentNetworkGuardrailUpdated Activity = 133 + // AgentNetworkGuardrailDeleted indicates that a user deleted an Agent Network guardrail + AgentNetworkGuardrailDeleted Activity = 134 + + // AgentNetworkBudgetRuleCreated indicates that a user created an Agent Network budget rule + AgentNetworkBudgetRuleCreated Activity = 135 + // AgentNetworkBudgetRuleUpdated indicates that a user updated an Agent Network budget rule + AgentNetworkBudgetRuleUpdated Activity = 136 + // AgentNetworkBudgetRuleDeleted indicates that a user deleted an Agent Network budget rule + AgentNetworkBudgetRuleDeleted Activity = 137 + + // AgentNetworkSettingsUpdated indicates that a user updated Agent Network account settings + AgentNetworkSettingsUpdated Activity = 139 + AccountDeleted Activity = 99999 ) @@ -400,6 +431,24 @@ var activityMap = map[Activity]Code{ UserExtendedPeerSession: {"User extended peer session", "user.peer.session.extend"}, + AgentNetworkProviderCreated: {"Agent Network provider created", "agent_network.provider.create"}, + AgentNetworkProviderUpdated: {"Agent Network provider updated", "agent_network.provider.update"}, + AgentNetworkProviderDeleted: {"Agent Network provider deleted", "agent_network.provider.delete"}, + + AgentNetworkPolicyCreated: {"Agent Network policy created", "agent_network.policy.create"}, + AgentNetworkPolicyUpdated: {"Agent Network policy updated", "agent_network.policy.update"}, + AgentNetworkPolicyDeleted: {"Agent Network policy deleted", "agent_network.policy.delete"}, + + AgentNetworkGuardrailCreated: {"Agent Network guardrail created", "agent_network.guardrail.create"}, + AgentNetworkGuardrailUpdated: {"Agent Network guardrail updated", "agent_network.guardrail.update"}, + AgentNetworkGuardrailDeleted: {"Agent Network guardrail deleted", "agent_network.guardrail.delete"}, + + AgentNetworkBudgetRuleCreated: {"Agent Network budget rule created", "agent_network.budget_rule.create"}, + AgentNetworkBudgetRuleUpdated: {"Agent Network budget rule updated", "agent_network.budget_rule.update"}, + AgentNetworkBudgetRuleDeleted: {"Agent Network budget rule deleted", "agent_network.budget_rule.delete"}, + + AgentNetworkSettingsUpdated: {"Agent Network settings updated", "agent_network.settings.update"}, + DomainAdded: {"Domain added", "domain.add"}, DomainDeleted: {"Domain deleted", "domain.delete"}, DomainValidated: {"Domain validated", "domain.validate"}, diff --git a/management/server/agentnetwork/catalog/catalog.go b/management/server/agentnetwork/catalog/catalog.go new file mode 100644 index 000000000..baf622778 --- /dev/null +++ b/management/server/agentnetwork/catalog/catalog.go @@ -0,0 +1,749 @@ +// Package catalog defines the static set of Agent Network providers +// recognized by the management server. The catalog is consulted both to +// validate provider_id on create/update and to surface the available +// providers (and their models) to the dashboard. +package catalog + +import "github.com/netbirdio/netbird/shared/management/http/api" + +// Model is the in-memory representation of a catalog model. +type Model struct { + ID string + Label string + InputPer1k float64 + OutputPer1k float64 + ContextWindow int +} + +// ProviderKind groups catalog entries for UI presentation. The split +// is semantic, not technical: +// - KindProvider: the upstream is a vendor's first-party API (OpenAI, +// Anthropic, Mistral, Bedrock, etc.) — NetBird talks straight to +// the model provider. +// - KindGateway: the upstream is itself a routing / aggregation layer +// in front of multiple providers (LiteLLM, Portkey, Helicone, …). +// These typically need NetBird identity stamped onto upstream +// requests so the gateway's analytics and budgets attribute to the +// real caller; that's what IdentityInjection is for. +// - KindCustom: the catch-all "OpenAI-compatible self-hosted endpoint" +// entry (vLLM, Ollama, custom inference servers). +// +// Frontend uses Kind to group the provider Select in the modal so an +// operator can spot at a glance which catalog entries proxy other +// providers vs. talk straight to one. Backend doesn't dispatch on Kind +// today; it's purely a presentation hint. +type ProviderKind string + +const ( + KindProvider ProviderKind = "provider" + KindGateway ProviderKind = "gateway" + KindCustom ProviderKind = "custom" +) + +// Provider is the in-memory representation of a catalog provider. +type Provider struct { + ID string + Name string + Description string + DefaultHost string + // Kind groups this entry for UI presentation; see ProviderKind. + Kind ProviderKind + // AuthHeaderName is the HTTP header the provider's API expects + // the credential under (e.g. "Authorization" for OpenAI, + // "x-api-key" for Anthropic). Combined with AuthHeaderTemplate + // at synthesis time to inject the auth header on every upstream + // request. + AuthHeaderName string + AuthHeaderTemplate string + DefaultContentType string + BrandColor string + // ParserID names the proxy LLM parser surface this provider + // speaks (matches llm.Parser.ProviderName: "openai", + // "anthropic"). Multiple catalog ids may share a parser surface + // (e.g. azure_openai_api and mistral_api both speak the OpenAI + // shape). Empty when no parser is yet implemented for the + // surface — the proxy middleware then falls back to URL sniffing + // or skips request-side enrichment. + ParserID string + // IdentityInjection, when non-nil, instructs the proxy to stamp + // the caller's NetBird identity onto upstream requests under the + // configured header names. Used for gateways like LiteLLM that + // key budgets and attribution off request headers (the gateway + // otherwise has no way to learn which user / group made the call). + // The proxy strips the same header names from the inbound request + // before stamping ours, so an app can't spoof identity by setting + // these headers itself. + IdentityInjection *IdentityInjection + // ExtraHeaders is a catalog-declared list of additional per- + // provider routing/config headers the proxy stamps on every + // upstream request. Distinct from AuthHeaderName/Template (which + // always carries the API_KEY) and from IdentityInjection (caller + // identity). Each entry surfaces an optional input on the + // dashboard's provider modal whose value lives on the provider + // record's ExtraValues map (keyed by ExtraHeader.Name). Empty + // list = no extra inputs rendered. Used today by Portkey for + // "x-portkey-config: pc-..." (a saved-config id that resolves + // upstream provider + credentials on Portkey's hosted side). + ExtraHeaders []ExtraHeader + Models []Model +} + +// ExtraHeader names a single optional per-provider routing/config +// header. Catalog declares N of these per provider type; the operator +// fills any subset on the provider record (see Provider.ExtraValues). +// At synth time, only entries with a non-empty operator value are +// stamped; the proxy's identity-inject middleware applies anti-spoof +// (Remove + Add) so a client can't supply these headers themselves. +// +// UI copy (label / help text / tooltip) for each known Name lives on +// the dashboard, not here — the backend's job is just to declare +// which wire headers are accepted. New provider needs an extra +// header? Add the Name here AND the matching UI copy on the dashboard. +type ExtraHeader struct { + // Name is the wire header name, e.g. "x-portkey-config". + Name string +} + +// IdentityInjection describes how the proxy stamps NetBird identity onto +// upstream gateway requests. Exactly one shape must be set — they're +// mutually exclusive and dispatched by the inject middleware. +// +// Shape choice tracks the wire convention the upstream gateway uses, +// not the vendor name. New gateways with a known shape become a catalog +// entry, not a new code path. +type IdentityInjection struct { + // HeaderPair emits separate headers per identity dimension + // (end-user id, tags as CSV). LiteLLM and OpenAI-compatible + // self-hosted gateways that read identity from dedicated headers. + HeaderPair *HeaderPairInjection + // JSONMetadata emits a single header carrying a JSON object with + // reserved keys for user / groups / etc. Portkey, Helicone-style + // metadata headers, anything that wants a structured envelope. + JSONMetadata *JSONMetadataInjection +} + +// HeaderPairInjection is the LiteLLM-style wire convention. +type HeaderPairInjection struct { + // Customizable, when true, marks the wire header names as + // operator-overridable: the dashboard surfaces EndUserIDHeader + // and TagsHeader as editable inputs (defaults shown as + // placeholders) and the synthesizer pulls the actual values from + // the provider record's IdentityHeader* fields rather than from + // these defaults. An empty operator value disables stamping for + // that dimension. Used today for Bifrost, whose log-metadata / + // telemetry header prefix (x-bf-lh-* vs x-bf-dim-*) is a + // per-operator choice; LiteLLM and similar gateways with a fixed + // wire protocol leave this false so the catalog defaults are + // authoritative. + Customizable bool + // EndUserIDHeader receives the caller's display identity (user + // email when the peer is attached to a user, else peer.Name), + // e.g. "x-litellm-end-user-id". + EndUserIDHeader string + // TagsHeader receives the caller's NetBird group display names + // as a CSV, e.g. "x-litellm-tags". + TagsHeader string + // TagsInBody, when true, additionally writes the tag list into + // the request body's metadata.tags array (a JSON path the + // gateway parses for budget enforcement). LiteLLM only honours + // metadata.tags for tag-budget gating — its x-litellm-tags + // header path feeds spend tracking but bypasses + // _tag_max_budget_check entirely. Body inject is skipped when + // the request body is empty, truncated, non-JSON, or when an + // existing metadata field is a non-object value (defensive: we + // never clobber a client-supplied non-object). The header path + // remains a robust fallback for spend tracking in those cases. + TagsInBody bool + // EndUserIDInBody, when true, additionally writes the display + // identity into the request body's top-level "user" field (the + // OpenAI-standard end-user identifier). LiteLLM resolves the end + // user id from headers first then body, so for LiteLLM this is + // belt-and-suspenders. It matters when an OpenAI-compatible + // gateway downstream of LiteLLM (or OpenAI direct, bypassing + // LiteLLM) only reads the body, and as anti-spoof: client- + // supplied "user" values are overwritten with our trusted + // identity. Same skip rules as TagsInBody. + EndUserIDInBody bool +} + +// JSONMetadataInjection is the Portkey-style wire convention: a single +// header carrying a JSON object. NetBird identity fields land under the +// configured reserved keys; missing keys (empty string) are skipped at +// emit time. +type JSONMetadataInjection struct { + // Customizable, when true, marks the JSON keys as operator- + // overridable. The dashboard surfaces UserKey and GroupsKey as + // editable inputs (the catalog values shown as placeholders) and + // the synthesizer pulls the actual JSON-key names from the + // provider record's IdentityHeader* fields. Same field reuse as + // HeaderPair's customizable path — the dimensions (user identity, + // groups) are the same, only the wire encoding differs (JSON key + // vs HTTP header name). An empty operator value disables emission + // for that dimension. Used today for Cloudflare AI Gateway, whose + // cf-aig-metadata header accepts arbitrary JSON keys; Portkey + // leaves this false because its keys are reserved by the Portkey + // schema. + Customizable bool + // Header is the wire header name carrying the JSON payload, e.g. + // "x-portkey-metadata". + Header string + // UserKey is the JSON key for the caller's display identity. + // Portkey reserves "_user" for this dimension. + UserKey string + // GroupsKey is the JSON key for the caller's NetBird groups, + // emitted as a CSV string value (Portkey requires string values). + GroupsKey string + // MaxValueLength caps each emitted JSON value, in bytes. Portkey + // enforces a 128-char limit per value; oversized values are + // truncated rather than failing the request. 0 disables the cap. + MaxValueLength int +} + +// providers is the canonical list of supported Agent Network providers. +// Update this list together with the dashboard's PROVIDER_CATALOG. +var providers = []Provider{ + { + ID: "openai_api", + Kind: KindProvider, + Name: "OpenAI API", + Description: "GPT, Responses API, and Embeddings", + DefaultHost: "api.openai.com", + AuthHeaderName: "Authorization", + AuthHeaderTemplate: "Bearer ${API_KEY}", + DefaultContentType: "application/json", + BrandColor: "#10A37F", + ParserID: "openai", + // Pricing + context windows cross-checked against LiteLLM's + // model_prices_and_context_window.json. Notable corrections from + // earlier values: o4-mini repriced from $4/$16 to $1.10/$4.40 + // per MTok, gpt-4o from $5/$15 to $2.50/$10, and the GPT-5 + // family context windows split between 1.05M for full-size + // models and 272K for mini/nano/codex variants. + Models: []Model{ + {ID: "gpt-5.5", Label: "GPT-5.5", InputPer1k: 0.005, OutputPer1k: 0.030, ContextWindow: 1050000}, + {ID: "gpt-5.5-pro", Label: "GPT-5.5 Pro", InputPer1k: 0.030, OutputPer1k: 0.180, ContextWindow: 1050000}, + {ID: "gpt-5.4", Label: "GPT-5.4", InputPer1k: 0.0025, OutputPer1k: 0.015, ContextWindow: 1050000}, + {ID: "gpt-5.4-pro", Label: "GPT-5.4 Pro", InputPer1k: 0.030, OutputPer1k: 0.180, ContextWindow: 1050000}, + {ID: "gpt-5.4-mini", Label: "GPT-5.4 Mini", InputPer1k: 0.00075, OutputPer1k: 0.0045, ContextWindow: 272000}, + {ID: "gpt-5.4-nano", Label: "GPT-5.4 Nano", InputPer1k: 0.0002, OutputPer1k: 0.00125, ContextWindow: 272000}, + {ID: "gpt-5.3-codex", Label: "GPT-5.3 Codex", InputPer1k: 0.00175, OutputPer1k: 0.014, ContextWindow: 272000}, + {ID: "gpt-5.3-chat-latest", Label: "GPT-5.3 Chat", InputPer1k: 0.00175, OutputPer1k: 0.014, ContextWindow: 128000}, + {ID: "o4-mini", Label: "o4-mini", InputPer1k: 0.0011, OutputPer1k: 0.0044, ContextWindow: 200000}, + {ID: "gpt-4.1", Label: "GPT-4.1", InputPer1k: 0.002, OutputPer1k: 0.008, ContextWindow: 1047576}, + {ID: "gpt-4.1-mini", Label: "GPT-4.1 mini", InputPer1k: 0.0004, OutputPer1k: 0.0016, ContextWindow: 1047576}, + {ID: "gpt-4.1-nano", Label: "GPT-4.1 nano", InputPer1k: 0.0001, OutputPer1k: 0.0004, ContextWindow: 1047576}, + {ID: "gpt-4o", Label: "GPT-4o", InputPer1k: 0.0025, OutputPer1k: 0.010, ContextWindow: 128000}, + {ID: "gpt-4o-mini", Label: "GPT-4o mini", InputPer1k: 0.00015, OutputPer1k: 0.0006, ContextWindow: 128000}, + {ID: "gpt-4-turbo", Label: "GPT-4 Turbo", InputPer1k: 0.01, OutputPer1k: 0.03, ContextWindow: 128000}, + {ID: "gpt-3.5-turbo", Label: "GPT-3.5 Turbo", InputPer1k: 0.0005, OutputPer1k: 0.0015, ContextWindow: 16385}, + {ID: "text-embedding-3-large", Label: "text-embedding-3-large", InputPer1k: 0.00013, OutputPer1k: 0, ContextWindow: 8191}, + {ID: "text-embedding-3-small", Label: "text-embedding-3-small", InputPer1k: 0.00002, OutputPer1k: 0, ContextWindow: 8191}, + }, + }, + { + ID: "anthropic_api", + Kind: KindProvider, + Name: "Anthropic API", + Description: "Claude Messages API", + DefaultHost: "api.anthropic.com", + AuthHeaderName: "x-api-key", + AuthHeaderTemplate: "${API_KEY}", + DefaultContentType: "application/json", + BrandColor: "#D97757", + ParserID: "anthropic", + // Per Anthropic's current model lineup. Pricing in USD per 1k + // tokens. Context windows: 4.6+ family is 1M; Haiku 4.5 stays at + // 200K. claude-3-7-sonnet and claude-3-5-haiku retired + // 2026-02-19 — dropped from the catalog. claude-opus-4-1 + // deprecated, retires 2026-08-05 — kept until the cutover. + // claude-mythos-5 omitted: Project Glasswing access only, not a + // general-availability target. claude-fable-5 requires the + // account to be on >= 30-day data retention or all requests + // 400. + Models: []Model{ + {ID: "claude-fable-5", Label: "Claude Fable 5", InputPer1k: 0.010, OutputPer1k: 0.050, ContextWindow: 1000000}, + {ID: "claude-opus-4-8", Label: "Claude Opus 4.8", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000}, + {ID: "claude-opus-4-7", Label: "Claude Opus 4.7", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000}, + {ID: "claude-opus-4-6", Label: "Claude Opus 4.6", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000}, + {ID: "claude-opus-4-1", Label: "Claude Opus 4.1 (deprecated, retires 2026-08-05)", InputPer1k: 0.015, OutputPer1k: 0.075, ContextWindow: 200000}, + {ID: "claude-sonnet-4-6", Label: "Claude Sonnet 4.6", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 1000000}, + {ID: "claude-sonnet-4-5", Label: "Claude Sonnet 4.5", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 200000}, + {ID: "claude-haiku-4-5", Label: "Claude Haiku 4.5", InputPer1k: 0.001, OutputPer1k: 0.005, ContextWindow: 200000}, + }, + }, + { + ID: "azure_openai_api", + Kind: KindProvider, + Name: "Azure OpenAI API", + Description: "Azure-hosted OpenAI deployments", + DefaultHost: ".openai.azure.com", + AuthHeaderName: "api-key", + AuthHeaderTemplate: "${API_KEY}", + DefaultContentType: "application/json", + BrandColor: "#0078D4", + ParserID: "openai", + // Mirrors openai_api pricing — Azure resells OpenAI models at the + // same per-token rates, just under different deployment names. + Models: []Model{ + {ID: "gpt-5.5", Label: "GPT-5.5 (Azure)", InputPer1k: 0.005, OutputPer1k: 0.030, ContextWindow: 1050000}, + {ID: "gpt-5.4", Label: "GPT-5.4 (Azure)", InputPer1k: 0.0025, OutputPer1k: 0.015, ContextWindow: 1050000}, + {ID: "gpt-5.4-mini", Label: "GPT-5.4 Mini (Azure)", InputPer1k: 0.00075, OutputPer1k: 0.0045, ContextWindow: 272000}, + {ID: "gpt-5.4-nano", Label: "GPT-5.4 Nano (Azure)", InputPer1k: 0.0002, OutputPer1k: 0.00125, ContextWindow: 272000}, + {ID: "o4-mini", Label: "o4-mini (Azure)", InputPer1k: 0.0011, OutputPer1k: 0.0044, ContextWindow: 200000}, + {ID: "gpt-4.1", Label: "GPT-4.1 (Azure)", InputPer1k: 0.002, OutputPer1k: 0.008, ContextWindow: 1047576}, + {ID: "gpt-4.1-mini", Label: "GPT-4.1 mini (Azure)", InputPer1k: 0.0004, OutputPer1k: 0.0016, ContextWindow: 1047576}, + {ID: "gpt-4o", Label: "GPT-4o (Azure)", InputPer1k: 0.0025, OutputPer1k: 0.010, ContextWindow: 128000}, + {ID: "gpt-4o-mini", Label: "GPT-4o mini (Azure)", InputPer1k: 0.00015, OutputPer1k: 0.0006, ContextWindow: 128000}, + {ID: "gpt-35-turbo", Label: "GPT-3.5 Turbo (Azure)", InputPer1k: 0.0005, OutputPer1k: 0.0015, ContextWindow: 16385}, + }, + }, + { + ID: "bedrock_api", + Kind: KindProvider, + Name: "AWS Bedrock API", + Description: "Anthropic, Meta, Cohere via Bedrock", + DefaultHost: "bedrock-runtime..amazonaws.com", + AuthHeaderName: "Authorization", + AuthHeaderTemplate: "Bearer ${API_KEY}", + DefaultContentType: "application/json", + BrandColor: "#FF9900", + // Anthropic models on Bedrock take the anthropic.* prefix and + // follow the same lineup / pricing as the first-party Anthropic + // catalog entry above. claude-3-7-sonnet and claude-3-5-haiku + // were retired upstream on 2026-02-19 — dropped from the + // Bedrock list too. Amazon Nova entries cross-checked against + // LiteLLM (added Nova Micro + the new Nova 2 Lite preview). + // Llama 3.3 70B entry kept unchanged — LiteLLM tracks only + // per-region Llama 3 entries; standalone 3.3 not yet listed. + Models: []Model{ + {ID: "anthropic.claude-opus-4-8", Label: "Claude Opus 4.8 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000}, + {ID: "anthropic.claude-opus-4-7", Label: "Claude Opus 4.7 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000}, + {ID: "anthropic.claude-opus-4-6", Label: "Claude Opus 4.6 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000}, + {ID: "anthropic.claude-opus-4-1", Label: "Claude Opus 4.1 (Bedrock, deprecated 2026-08-05)", InputPer1k: 0.015, OutputPer1k: 0.075, ContextWindow: 200000}, + {ID: "anthropic.claude-sonnet-4-6", Label: "Claude Sonnet 4.6 (Bedrock)", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 1000000}, + {ID: "anthropic.claude-sonnet-4-5", Label: "Claude Sonnet 4.5 (Bedrock)", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 200000}, + {ID: "anthropic.claude-haiku-4-5", Label: "Claude Haiku 4.5 (Bedrock)", InputPer1k: 0.001, OutputPer1k: 0.005, ContextWindow: 200000}, + {ID: "meta.llama3-3-70b-instruct", Label: "Llama 3.3 70B (Bedrock)", InputPer1k: 0.00072, OutputPer1k: 0.00072, ContextWindow: 128000}, + {ID: "amazon.nova-2-lite", Label: "Amazon Nova 2 Lite (Bedrock, preview)", InputPer1k: 0.0003, OutputPer1k: 0.0025, ContextWindow: 1000000}, + {ID: "amazon.nova-pro", Label: "Amazon Nova Pro (Bedrock)", InputPer1k: 0.0008, OutputPer1k: 0.0032, ContextWindow: 300000}, + {ID: "amazon.nova-lite", Label: "Amazon Nova Lite (Bedrock)", InputPer1k: 0.00006, OutputPer1k: 0.00024, ContextWindow: 300000}, + {ID: "amazon.nova-micro", Label: "Amazon Nova Micro (Bedrock)", InputPer1k: 0.000035, OutputPer1k: 0.00014, ContextWindow: 128000}, + }, + }, + { + ID: "vertex_ai_api", + Kind: KindProvider, + Name: "Google Vertex AI API", + Description: "Anthropic Claude models hosted on Vertex AI", + DefaultHost: "-aiplatform.googleapis.com", + AuthHeaderName: "Authorization", + AuthHeaderTemplate: "Bearer ${API_KEY}", + DefaultContentType: "application/json", + BrandColor: "#4285F4", + // Vertex carries the model in the URL path and authenticates with a + // service-account-minted OAuth token (api_key = "keyfile::"). + // Only Anthropic-on-Vertex is metered today: the request parser maps the + // anthropic publisher to the Anthropic parser, so the lineup + prices + // mirror the first-party Anthropic catalog (LiteLLM vertex_ai/claude-* + // confirms the same per-token rates; cross-region profiles in eu/apac + // carry a ~10% premium that base pricing does not model). Gemini (the + // google publisher) is intentionally omitted until a Gemini parser + // exists — the router denies unmeterable publishers rather than forward + // them uncounted. + Models: []Model{ + {ID: "claude-fable-5", Label: "Claude Fable 5 (Vertex)", InputPer1k: 0.010, OutputPer1k: 0.050, ContextWindow: 1000000}, + {ID: "claude-opus-4-8", Label: "Claude Opus 4.8 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000}, + {ID: "claude-opus-4-7", Label: "Claude Opus 4.7 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000}, + {ID: "claude-opus-4-6", Label: "Claude Opus 4.6 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000}, + {ID: "claude-opus-4-1", Label: "Claude Opus 4.1 (Vertex, deprecated 2026-08-05)", InputPer1k: 0.015, OutputPer1k: 0.075, ContextWindow: 200000}, + {ID: "claude-sonnet-4-6", Label: "Claude Sonnet 4.6 (Vertex)", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 1000000}, + {ID: "claude-sonnet-4-5", Label: "Claude Sonnet 4.5 (Vertex)", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 200000}, + {ID: "claude-haiku-4-5", Label: "Claude Haiku 4.5 (Vertex)", InputPer1k: 0.001, OutputPer1k: 0.005, ContextWindow: 200000}, + }, + }, + { + ID: "mistral_api", + Kind: KindProvider, + Name: "Mistral API", + Description: "Mistral cloud API", + DefaultHost: "api.mistral.ai", + AuthHeaderName: "Authorization", + AuthHeaderTemplate: "Bearer ${API_KEY}", + DefaultContentType: "application/json", + BrandColor: "#FF7000", + ParserID: "openai", + // Pricing + context windows cross-checked against LiteLLM. Key + // gotchas the marketing page hides: + // - `mistral-medium-latest` aliases to Medium 3.1 ($0.40/$2), + // NOT Medium 3.5 ($1.50/$7.50). Catalog exposes both. + // - `mistral-large-latest` aliases to Large 3 — 262K context, + // cheaper than Medium 3.5. + // - Magistral models are tuned for reasoning but cap context + // at only 40K (vs 128K-262K elsewhere). + // - `codestral-latest` still routes to the old 2405 build + // ($1/$3) per LiteLLM; the newer codestral-2508 is both + // cheaper and longer-context. Both exposed. + // - Pixtral was folded into the main Large/Medium series; no + // standalone vision entry. + Models: []Model{ + {ID: "mistral-large-latest", Label: "Mistral Large 3", InputPer1k: 0.0005, OutputPer1k: 0.0015, ContextWindow: 262144}, + {ID: "mistral-medium-latest", Label: "Mistral Medium 3.1", InputPer1k: 0.0004, OutputPer1k: 0.002, ContextWindow: 131072}, + {ID: "mistral-medium-3-5", Label: "Mistral Medium 3.5", InputPer1k: 0.0015, OutputPer1k: 0.0075, ContextWindow: 262144}, + {ID: "mistral-small-latest", Label: "Mistral Small 3.2", InputPer1k: 0.00006, OutputPer1k: 0.00018, ContextWindow: 131072}, + {ID: "magistral-medium-latest", Label: "Magistral Medium (reasoning)", InputPer1k: 0.002, OutputPer1k: 0.005, ContextWindow: 40000}, + {ID: "magistral-small-latest", Label: "Magistral Small (reasoning)", InputPer1k: 0.0005, OutputPer1k: 0.0015, ContextWindow: 40000}, + {ID: "devstral-medium-latest", Label: "Devstral Medium 2 (coding)", InputPer1k: 0.0004, OutputPer1k: 0.002, ContextWindow: 256000}, + {ID: "devstral-small-latest", Label: "Devstral Small 2 (coding)", InputPer1k: 0.0001, OutputPer1k: 0.0003, ContextWindow: 256000}, + {ID: "codestral-2508", Label: "Codestral 2508", InputPer1k: 0.0003, OutputPer1k: 0.0009, ContextWindow: 256000}, + {ID: "codestral-latest", Label: "Codestral (legacy 2405)", InputPer1k: 0.001, OutputPer1k: 0.003, ContextWindow: 32000}, + {ID: "ministral-3-14b-2512", Label: "Ministral 3 14B", InputPer1k: 0.0002, OutputPer1k: 0.0002, ContextWindow: 262144}, + {ID: "ministral-8b-latest", Label: "Ministral 8B", InputPer1k: 0.00015, OutputPer1k: 0.00015, ContextWindow: 262144}, + {ID: "ministral-3-3b-2512", Label: "Ministral 3 3B", InputPer1k: 0.0001, OutputPer1k: 0.0001, ContextWindow: 131072}, + {ID: "mistral-embed", Label: "Mistral Embed", InputPer1k: 0.0001, OutputPer1k: 0, ContextWindow: 8192}, + }, + }, + { + ID: "litellm_proxy", + Kind: KindGateway, + Name: "LiteLLM Proxy", + Description: "Bring your own LiteLLM proxy with NetBird identity stamped on every request", + DefaultHost: "", + AuthHeaderName: "Authorization", + AuthHeaderTemplate: "Bearer ${API_KEY}", + DefaultContentType: "application/json", + BrandColor: "#0EA5E9", + ParserID: "openai", + // IdentityInjection requires a LiteLLM virtual key minted with + // metadata.allow_client_tags=true; the master key silently drops + // caller tags. Tags go out via both the x-litellm-tags header and + // body metadata.tags: LiteLLM enforces budgets from the body only, + // so the header is the spend-tracking fallback when body injection + // can't run. See the Agent Network provider docs for key setup. + IdentityInjection: &IdentityInjection{ + HeaderPair: &HeaderPairInjection{ + EndUserIDHeader: "x-litellm-end-user-id", + TagsHeader: "x-litellm-tags", + TagsInBody: true, + EndUserIDInBody: true, + }, + }, + Models: []Model{}, + }, + { + ID: "portkey", + Kind: KindGateway, + Name: "Portkey AI Gateway", + Description: "Portkey AI Gateway with NetBird identity stamped via x-portkey-metadata", + DefaultHost: "api.portkey.ai", + // Portkey hosted requires x-portkey-api-key (account key) + // plus a routing decision per request. The simplest routing + // path is a saved Portkey config id stamped via + // x-portkey-config — operators paste the pc-... id once and + // Portkey resolves the upstream provider + virtual key from + // it. ExtraHeaders below surfaces the input. Alternative: + // callers author "@org/model" in the body; both flows + // coexist (per-request authoring still works without a + // configured value). + AuthHeaderName: "x-portkey-api-key", + AuthHeaderTemplate: "${API_KEY}", + DefaultContentType: "application/json", + BrandColor: "#FF5C00", + ParserID: "openai", + IdentityInjection: &IdentityInjection{ + JSONMetadata: &JSONMetadataInjection{ + Header: "x-portkey-metadata", + UserKey: "_user", + GroupsKey: "groups", + MaxValueLength: 128, + }, + }, + ExtraHeaders: []ExtraHeader{ + {Name: "x-portkey-config"}, + }, + Models: []Model{}, + }, + { + ID: "bifrost", + Kind: KindGateway, + Name: "Bifrost", + Description: "Maxim AI's Bifrost gateway. Point upstream URL at /openai/v1 or /anthropic/v1 on your Bifrost host depending on which body shape your apps use.", + DefaultHost: "", + AuthHeaderName: "Authorization", + AuthHeaderTemplate: "Bearer ${API_KEY}", + DefaultContentType: "application/json", + BrandColor: "#7C3AED", + // ParserID empty: the proxy's request parser sniffs the URL + // path. Bifrost's /openai/v1/... contains "/v1/chat/completions" + // (matches OpenAIParser.DetectFromURL); /anthropic/v1/messages + // contains "/v1/messages" (matches AnthropicParser). Operators + // who paste a different prefix get no usage parsing and the + // cost meter skips with skipMissingProvider — degraded but + // non-fatal. + ParserID: "", + // Identity-injection headers are operator-customisable. The + // HeaderPair values below are PLACEHOLDERS surfaced by the + // dashboard; the actual values stamped on the wire come from + // the provider record's IdentityHeaderUserID / + // IdentityHeaderGroups fields. An empty operator value + // disables stamping for that dimension (the inject middleware + // already no-ops on empty header names). Defaulting to the + // x-bf-dim- family so the values land in Bifrost's + // Prometheus/OTEL pipelines when the operator declares the + // label names in their client.prometheus_labels config — see + // docs.getbifrost.ai/features/telemetry. Operators who use + // the always-on x-bf-lh- log-metadata family (no Bifrost-side + // declaration required) just edit the inputs. + // + // Bifrost virtual keys (sk-bf-*) ride Authorization: Bearer. + // Operators provision the VK on their Bifrost (UI / + // config.json / POST /api/governance/virtual-keys) and paste + // the returned sk-bf-... as ${API_KEY}. Pin v1.4+ to avoid + // the v1.3.0 x-bf-vk regression (maximhq/bifrost#632). + IdentityInjection: &IdentityInjection{ + HeaderPair: &HeaderPairInjection{ + EndUserIDHeader: "x-bf-dim-netbird_user_id", + TagsHeader: "x-bf-dim-netbird_groups", + Customizable: true, + }, + }, + Models: []Model{}, + }, + { + ID: "cloudflare_ai_gateway", + Kind: KindGateway, + Name: "Cloudflare AI Gateway", + Description: "Cloudflare AI Gateway. Operator pastes the gateway URL (with the upstream provider slug like /openai or /anthropic so the URL sniffer dispatches to the right parser) and a per-gateway authentication token. Recommended setup is BYOK / Stored Keys: Cloudflare manages the upstream provider credential and the gateway token is the only secret NetBird needs.", + DefaultHost: "", + AuthHeaderName: "cf-aig-authorization", + AuthHeaderTemplate: "Bearer ${API_KEY}", + DefaultContentType: "application/json", + BrandColor: "#F38020", + // ParserID empty: like Bifrost, the proxy's parser-detect + // sniffs the URL path. /openai/... contains the OpenAI hint + // substrings; /anthropic/v1/messages contains /v1/messages + // (matches AnthropicParser). The /compat universal endpoint + // also speaks OpenAI shape so OpenAIParser handles it. + // Operators who paste a different prefix degrade to no-cost + // (skipMissingProvider) but the request still flows. + ParserID: "", + // cf-aig-metadata is a single header carrying a JSON object; + // up to five string/number/boolean values per request. NetBird + // occupies two slots (user id + groups CSV) and leaves three + // for operator-added context. JSON keys are operator- + // customisable so Cloudflare-side log filters can use the + // operator's existing label conventions instead of NetBird's + // defaults — hence Customizable=true. The dashboard surfaces + // the catalog values as placeholders; only the values stored + // on the provider record's IdentityHeader* fields land on the + // wire (empty operator value = key is omitted from the JSON, + // since applyJSONMetadata already skips empty keys). + IdentityInjection: &IdentityInjection{ + JSONMetadata: &JSONMetadataInjection{ + Header: "cf-aig-metadata", + UserKey: "netbird_user_id", + GroupsKey: "netbird_groups", + Customizable: true, + // Cloudflare's docs don't specify a per-value cap; + // leaving 0 disables the truncate path. Header-level + // constraint is "5 entries max" rather than length. + MaxValueLength: 0, + }, + }, + Models: []Model{}, + }, + { + ID: "vercel_ai_gateway", + Kind: KindGateway, + Name: "Vercel AI Gateway", + Description: "Vercel's unified API for hundreds of models. Single endpoint, OpenAI-compatible body, model dispatch via prefix (openai/..., anthropic/..., google/..., xai/...). Per-user / per-tag attribution lands in Vercel's Custom Reporting API and observability dashboard.", + DefaultHost: "", + AuthHeaderName: "Authorization", + AuthHeaderTemplate: "Bearer ${API_KEY}", + DefaultContentType: "application/json", + BrandColor: "#000000", + // Vercel always speaks OpenAI shape on /v1/chat/completions — + // the model prefix in the body picks the upstream provider. + // No URL sniffing needed; pin the parser directly. + ParserID: "openai", + // HeaderPair shape with fixed wire names dictated by Vercel's + // Custom Reporting API contract. Customizable=false because + // renaming the headers makes Vercel silently stop attributing + // — the gateway's reporting endpoint only matches its own + // header names. Same fixed-protocol position as LiteLLM. + // + // Caveats operators should know: + // - up to 10 tags total per request (deduped); 11+ → HTTP 400 + // - each tag must be 1-64 chars + // - user up to 256 chars (NetBird user emails fit) + // - $0.075 per 1k unique user/tag values written + // We don't enforce the caps in the inject middleware today; + // operators in groups beyond the 10-tag limit will see Vercel + // 400s and need to re-scope their group memberships. + IdentityInjection: &IdentityInjection{ + HeaderPair: &HeaderPairInjection{ + EndUserIDHeader: "ai-reporting-user", + TagsHeader: "ai-reporting-tags", + }, + }, + Models: []Model{}, + }, + { + ID: "openrouter", + Kind: KindGateway, + Name: "OpenRouter", + Description: "OpenRouter's unified API for hundreds of models. Single endpoint at openrouter.ai/api/v1, OpenAI-compatible body, model dispatch via prefix (anthropic/claude-..., openai/gpt-..., google/gemini-..., etc.). Per-user attribution lands in OpenRouter's analytics via the OpenAI-standard `user` body field; OpenRouter has no groups / tags dimension at request time.", + DefaultHost: "openrouter.ai/api/v1", + AuthHeaderName: "Authorization", + AuthHeaderTemplate: "Bearer ${API_KEY}", + DefaultContentType: "application/json", + BrandColor: "#6F4FF2", + // OpenRouter is single-endpoint OpenAI-shape on /api/v1/chat/completions — + // model prefix in the body picks the upstream provider. + // Pinning the parser saves URL sniffing. + ParserID: "openai", + // HeaderPair shape with EndUserIDInBody as the only active + // dimension. OpenRouter's per-user attribution is the + // OpenAI-standard `user` body field, not a header — and + // OpenRouter offers no per-request groups / tags dimension at + // all. Customizable=false because the field name is locked by + // OpenAI's spec; renaming would just defeat the inject. + IdentityInjection: &IdentityInjection{ + HeaderPair: &HeaderPairInjection{ + EndUserIDInBody: true, + }, + }, + // HTTP-Referer + X-OpenRouter-Title surface in OpenRouter's + // app rankings and per-app analytics. Operators paste their + // own app URL + display name on the provider record so their + // requests show under their brand instead of "no app". Both + // are static per-deployment, not per-request, hence the + // ExtraHeaders mechanism (operator-typed value, stamped on + // every request to this provider). Skip X-OpenRouter-Categories + // for now — the marketplace-categories dimension is + // niche-enough that we'd add it on demand. + ExtraHeaders: []ExtraHeader{ + {Name: "HTTP-Referer"}, + {Name: "X-OpenRouter-Title"}, + }, + Models: []Model{}, + }, + { + ID: "custom", + Kind: KindCustom, + Name: "Custom / Self-hosted", + Description: "OpenAI-compatible endpoint (vLLM, Ollama, …)", + DefaultHost: "", + AuthHeaderName: "Authorization", + AuthHeaderTemplate: "Bearer ${API_KEY}", + DefaultContentType: "application/json", + BrandColor: "#9CA3AF", + Models: []Model{}, + }, +} + +// All returns a copy of the full catalog. +func All() []Provider { + out := make([]Provider, len(providers)) + copy(out, providers) + return out +} + +// Lookup returns the catalog entry with the given id, if any. +func Lookup(id string) (Provider, bool) { + for _, p := range providers { + if p.ID == id { + return p, true + } + } + return Provider{}, false +} + +// IsKnown reports whether the given id refers to a catalog entry. +func IsKnown(id string) bool { + _, ok := Lookup(id) + return ok +} + +// IsVertexPathStyle reports whether a provider uses the Google Vertex AI +// request shape — the model is carried in the URL path +// (/v1/projects/{p}/locations/{r}/publishers/{pub}/models/{model}:{action}) +// rather than the body, so the proxy routes it by path instead of by model. +func IsVertexPathStyle(providerID string) bool { + return providerID == "vertex_ai_api" +} + +// IsBedrockPathStyle reports whether a provider uses the AWS Bedrock request +// shape — the model is carried in the URL path (/model/{modelId}/{action}, +// action being invoke, invoke-with-response-stream, converse, or +// converse-stream) rather than the body, so the proxy routes it by path. +func IsBedrockPathStyle(providerID string) bool { + return providerID == "bedrock_api" +} + +// ToAPIResponse renders a catalog provider as the API representation. +func (p Provider) ToAPIResponse() api.AgentNetworkCatalogProvider { + models := make([]api.AgentNetworkCatalogModel, 0, len(p.Models)) + for _, m := range p.Models { + models = append(models, api.AgentNetworkCatalogModel{ + Id: m.ID, + Label: m.Label, + InputPer1k: m.InputPer1k, + OutputPer1k: m.OutputPer1k, + ContextWindow: m.ContextWindow, + }) + } + kind := api.AgentNetworkCatalogProviderKindProvider + switch p.Kind { + case KindGateway: + kind = api.AgentNetworkCatalogProviderKindGateway + case KindCustom: + kind = api.AgentNetworkCatalogProviderKindCustom + } + resp := api.AgentNetworkCatalogProvider{ + Id: p.ID, + Name: p.Name, + Description: p.Description, + DefaultHost: p.DefaultHost, + Kind: kind, + AuthHeaderTemplate: p.AuthHeaderTemplate, + DefaultContentType: p.DefaultContentType, + BrandColor: p.BrandColor, + Models: models, + } + if len(p.ExtraHeaders) > 0 { + extras := make([]api.AgentNetworkCatalogExtraHeader, 0, len(p.ExtraHeaders)) + for _, h := range p.ExtraHeaders { + extras = append(extras, api.AgentNetworkCatalogExtraHeader{ + Name: h.Name, + }) + } + resp.ExtraHeaders = &extras + } + // Surface IdentityInjection so the dashboard can decide whether + // to render editable inputs vs. a read-only mappings strip per + // shape's customizable flag. HeaderPair (Bifrost) and + // JSONMetadata (Cloudflare, Portkey) are mutually exclusive on a + // given catalog entry; emit whichever shape is set. + if p.IdentityInjection != nil { + injection := &api.AgentNetworkCatalogIdentityInjection{} + if hp := p.IdentityInjection.HeaderPair; hp != nil { + injection.HeaderPair = &api.AgentNetworkCatalogHeaderPairInjection{ + Customizable: hp.Customizable, + EndUserIdHeader: hp.EndUserIDHeader, + TagsHeader: hp.TagsHeader, + } + } + if jm := p.IdentityInjection.JSONMetadata; jm != nil { + injection.JsonMetadata = &api.AgentNetworkCatalogJSONMetadataInjection{ + Customizable: jm.Customizable, + Header: jm.Header, + UserKey: jm.UserKey, + GroupsKey: jm.GroupsKey, + } + } + if injection.HeaderPair != nil || injection.JsonMetadata != nil { + resp.IdentityInjection = injection + } + } + return resp +} diff --git a/management/server/agentnetwork/labelgen/labelgen.go b/management/server/agentnetwork/labelgen/labelgen.go new file mode 100644 index 000000000..b45ff4ea8 --- /dev/null +++ b/management/server/agentnetwork/labelgen/labelgen.go @@ -0,0 +1,66 @@ +// Package labelgen produces DNS-safe Agent Network subdomain labels. +package labelgen + +import ( + "fmt" + "math/rand" + "sort" + "sync" +) + +// pickAttempts caps the random retries before falling back to the +// suffixed form. Eight is a soft compromise: with a near-empty taken +// set the very first pick almost always succeeds; when the wordlist is +// densely populated the fallback eventually fires anyway. +const pickAttempts = 8 + +var ( + dedupOnce sync.Once + uniqWords []string +) + +// uniqueWords returns the wordlist deduplicated and sorted for +// deterministic exhaustion behaviour. Lazy-built once per process. +func uniqueWords() []string { + dedupOnce.Do(func() { + seen := make(map[string]struct{}, len(words)) + uniqWords = make([]string, 0, len(words)) + for _, w := range words { + if _, ok := seen[w]; ok { + continue + } + seen[w] = struct{}{} + uniqWords = append(uniqWords, w) + } + sort.Strings(uniqWords) + }) + return uniqWords +} + +// PickUnique selects a label not already in `taken`. It tries up to +// pickAttempts random picks; on exhaustion it scans the deduplicated +// wordlist for any remaining free entry, and if none is left appends +// `-` to a deterministic word and returns. The caller +// is responsible for seeding rng (math/rand). +func PickUnique(rng *rand.Rand, taken map[string]struct{}, fallbackSuffix string) string { + pool := uniqueWords() + if len(pool) == 0 { + return fallbackSuffix + } + + for i := 0; i < pickAttempts; i++ { + w := pool[rng.Intn(len(pool))] + if _, ok := taken[w]; !ok { + return w + } + } + + for _, w := range pool { + if _, ok := taken[w]; !ok { + return w + } + } + + w := pool[rng.Intn(len(pool))] + return fmt.Sprintf("%s-%s", w, fallbackSuffix) +} diff --git a/management/server/agentnetwork/labelgen/labelgen_test.go b/management/server/agentnetwork/labelgen/labelgen_test.go new file mode 100644 index 000000000..f03a3501d --- /dev/null +++ b/management/server/agentnetwork/labelgen/labelgen_test.go @@ -0,0 +1,101 @@ +package labelgen + +import ( + "math/rand" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestPickUnique_DeterministicWithSeededRng locks the property the +// caller relies on: same seed + same taken set → same pick. Without +// that, the bootstrap flow can't reproduce a label across retries. +func TestPickUnique_DeterministicWithSeededRng(t *testing.T) { + taken := map[string]struct{}{} + + rngA := rand.New(rand.NewSource(42)) + rngB := rand.New(rand.NewSource(42)) + + a := PickUnique(rngA, taken, "abcd") + b := PickUnique(rngB, taken, "abcd") + + assert.Equal(t, a, b, "Same seed and taken set must produce identical pick") +} + +// TestPickUnique_AvoidsTakenWordsWhenMostAreReserved seeds taken with +// every word in the pool except a handful and confirms PickUnique +// finds one of the remaining free entries instead of returning the +// fallback form. +func TestPickUnique_AvoidsTakenWordsWhenMostAreReserved(t *testing.T) { + pool := uniqueWords() + require.NotEmpty(t, pool, "wordlist must be populated for the test to mean anything") + + free := map[string]struct{}{ + pool[0]: {}, + pool[len(pool)/2]: {}, + pool[len(pool)-1]: {}, + } + + taken := make(map[string]struct{}, len(pool)) + for _, w := range pool { + if _, ok := free[w]; ok { + continue + } + taken[w] = struct{}{} + } + + rng := rand.New(rand.NewSource(7)) + got := PickUnique(rng, taken, "abcd") + + _, isFree := free[got] + assert.True(t, isFree, "PickUnique must return one of the free words; got %q", got) + assert.NotContains(t, got, "-", "Free pick must not be the suffix fallback form") +} + +// TestPickUnique_FallsBackWhenAllReserved exhausts the pool and +// confirms PickUnique appends the supplied suffix instead of +// returning a duplicate. +func TestPickUnique_FallsBackWhenAllReserved(t *testing.T) { + pool := uniqueWords() + + taken := make(map[string]struct{}, len(pool)) + for _, w := range pool { + taken[w] = struct{}{} + } + + rng := rand.New(rand.NewSource(99)) + got := PickUnique(rng, taken, "abcd") + + assert.True(t, strings.HasSuffix(got, "-abcd"), "Exhausted pool must produce -; got %q", got) + + prefix := strings.TrimSuffix(got, "-abcd") + found := false + for _, w := range pool { + if w == prefix { + found = true + break + } + } + assert.True(t, found, "Fallback prefix must be drawn from the wordlist; got %q", prefix) +} + +// TestUniqueWords_DropsDuplicates guards against authoring slips in +// words.go: every entry must be unique and DNS-safe. +func TestUniqueWords_DropsDuplicates(t *testing.T) { + pool := uniqueWords() + seen := make(map[string]struct{}, len(pool)) + for _, w := range pool { + _, dup := seen[w] + assert.False(t, dup, "Duplicate entry %q in deduplicated pool", w) + seen[w] = struct{}{} + assert.GreaterOrEqual(t, len(w), 4, "Word %q is shorter than 4 chars", w) + assert.LessOrEqual(t, len(w), 12, "Word %q is longer than 12 chars", w) + for _, r := range w { + ok := r >= 'a' && r <= 'z' + assert.True(t, ok, "Word %q contains non-lowercase-ASCII rune %q", w, r) + } + } + assert.GreaterOrEqual(t, len(pool), 500, "Pool must contain at least 500 unique words") +} diff --git a/management/server/agentnetwork/labelgen/words.go b/management/server/agentnetwork/labelgen/words.go new file mode 100644 index 000000000..80ebf2f51 --- /dev/null +++ b/management/server/agentnetwork/labelgen/words.go @@ -0,0 +1,136 @@ +// Package labelgen produces DNS-safe Agent Network subdomain labels. +// +// The wordlist below is a curated subset drawn from public-domain +// nature / common-noun pools (e.g. EFF's diceware lists). Every entry +// is lowercase ASCII, 4–12 chars, no hyphens, no digits, and was +// hand-checked to avoid offensive, brand, or region-specific terms. +package labelgen + +// words is the pool PickUnique selects from. The slice is intentionally +// not sorted — random picks distribute across the list naturally. +var words = []string{ + "acorn", "adobe", "agate", "alder", "almond", "alpine", "amber", "amethyst", + "anchor", "antler", "apple", "apricot", "arcade", "arctic", "arrow", "ashen", + "aspen", "atlas", "atom", "aurora", "autumn", "azure", + "badger", "bamboo", "banana", "banjo", "barley", "barn", "basalt", "basil", + "basin", "bayou", "beach", "beacon", "beaver", "beech", "beetle", "berry", + "birch", "bison", "blossom", "blue", "bobcat", "bonsai", "boulder", "branch", + "brass", "breeze", "bridge", "bright", "brook", "broom", "brown", "buffalo", + "bumble", "burrow", "butter", "button", + "cabin", "cactus", "calm", "camel", "campfire", "canary", "candle", "canoe", + "canyon", "cardinal", "carrot", "cascade", "castle", "cedar", "celery", "cello", + "cement", "cherry", "chestnut", "chime", "cinnamon", "cinder", "citron", "clay", + "clear", "cliff", "clock", "cloud", "clover", "coast", "cobalt", "cobble", + "cocoa", "coffee", "comet", "compass", "copper", "coral", "corner", "cosmos", + "cotton", "cougar", "country", "coyote", "cove", "crane", "crater", "creek", + "crescent", "crimson", "crocus", "crystal", "cypress", + "daffodil", "dahlia", "daisy", "dawn", "deer", "delta", "denim", "desert", + "dewdrop", "diamond", "dolphin", "doodle", "dove", "dragon", "drift", "drop", + "dune", "dusk", "dusty", + "eagle", "earth", "echo", "elder", "elkhorn", "ember", "emerald", "emperor", + "evergreen", "evening", + "falcon", "fawn", "feather", "fern", "fiddle", "field", "fiesta", "finch", + "firepit", "firefly", "fjord", "flame", "flax", "fleece", "flint", "floral", + "flower", "flute", "foal", "foggy", "forest", "fountain", "foxglove", "fresh", + "frost", "fuchsia", "fudge", + "gable", "galaxy", "garden", "garnet", "gazelle", "geode", "geyser", "ginger", + "glacier", "glade", "glass", "glow", "gold", "goose", "gorge", "gourd", + "granite", "grape", "grass", "gravel", "grayling", "greenery", "grizzly", "grove", + "gull", "gumdrop", "gust", + "hammock", "harbor", "harvest", "hawk", "hazel", "heather", "hedge", "heron", + "hibiscus", "hickory", "hideaway", "highland", "hill", "hive", "hollow", "honey", + "hopper", "horizon", "hummingbird", "husky", + "iceberg", "indigo", "iris", "island", "ivory", "ivybush", + "jade", "jasmine", "jasper", "jaybird", "jelly", "jewel", "jonquil", "journey", + "juniper", "jupiter", "jute", + "kale", "kangaroo", "kayak", "kelp", "kestrel", "kettle", "khaki", "kindling", + "kingfisher", "kiwi", "knapweed", "koala", + "lagoon", "lake", "lantern", "larch", "lark", "laurel", "lava", "lavender", + "leaf", "lemon", "lichen", "light", "lilac", "lily", "lime", "limestone", + "linden", "linen", "lion", "lobster", "locust", "loon", "lotus", "lumber", + "lunar", "lupine", "lynx", + "madrone", "magenta", "magnolia", "mahogany", "mallow", "mango", "manor", "maple", + "marble", "marigold", "marina", "marlin", "marsh", "mauve", "meadow", "melody", + "melon", "merlin", "metal", "midnight", "milk", "millet", "mineral", "mint", + "mirror", "mist", "mitten", "molasses", "moon", "moose", "morning", "moss", + "mountain", "mulberry", "muscat", "mustard", + "narwhal", "navy", "nectar", "needle", "nest", "nettle", "newt", "nightfall", + "noon", "nook", "north", "nova", "nutmeg", + "oaken", "oasis", "oatmeal", "ocean", "ochre", "octagon", "olive", "onyx", + "opal", "orange", "orbit", "orchard", "orchid", "oregano", "orion", "osprey", + "otter", "outpost", "owlet", "oyster", + "painter", "palace", "palm", "pansy", "panther", "papaya", "paprika", "parsley", + "partridge", "passage", "pastel", "patio", "peach", "peacock", "pear", "pearl", + "pebble", "pecan", "pelican", "penguin", "peony", "pepper", "perch", "peridot", + "pewter", "phoenix", "pier", "pillar", "pine", "pineapple", "pinto", "piper", + "pistachio", "plain", "planet", "plateau", "platinum", "plum", "plume", "polar", + "pollen", "pond", "poplar", "poppy", "porcelain", "portal", "portrait", "potato", + "prairie", "primrose", "prism", "puffin", "pumpkin", + "quail", "quartz", "quaver", "quill", "quince", "quinoa", + "rabbit", "raccoon", "radish", "rain", "rainbow", "raindrop", "rapids", "raspberry", + "raven", "ravine", "redwood", "reed", "reef", "ridge", "river", "robin", + "rocket", "rubyred", "rose", "rosemary", "rosewood", "ruffle", "rugby", "russet", + "rustic", "ryefield", + "saffron", "sage", "salmon", "sand", "sandstone", "sapphire", "savanna", "scarlet", + "scout", "seal", "season", "seaweed", "sequoia", "shadow", "shamrock", "shell", + "sherbet", "shore", "silver", "siskin", "skybloom", "skyline", "sleet", "smoke", + "snail", "snapdragon", "snow", "snowflake", "snowy", "solar", "song", "sonic", + "sorrel", "south", "sparkle", "sparrow", "spice", "spider", "spinach", "spire", + "spring", "sprout", "spruce", "squirrel", "starfish", "starlight", "stoat", "stone", + "stork", "storm", "stream", "studio", "summer", "sunbeam", "sundew", "sunny", + "sunrise", "sunset", "swallow", "swan", "sweet", "sycamore", + "tangelo", "tangerine", "tansy", "taupe", "teak", "teal", "thicket", "thistle", + "thrush", "thunder", "tide", "tiger", "tinder", "topaz", "torch", "tortoise", + "tower", "trail", "tranquil", "tundra", "tulip", "turquoise", "turtle", "twig", + "twilight", + "umber", "uplands", + "valley", "vanilla", "velvet", "venus", "verdant", "verdigris", "vermillion", "violet", + "vista", "vivid", "volcano", "vortex", + "walnut", "warbler", "watercress", "waterfall", "wave", "waxwing", "weasel", "westwind", + "whale", "whisker", "whisper", "wicker", "wildwood", "willow", "winter", "wisp", + "wisteria", "wolf", "wombat", "woodland", "woolly", "wren", "wreath", + "yarrow", "yellow", "yewtree", "yodel", + "zebra", "zenith", "zephyr", "zinnia", + "alabaster", "alfalfa", "almanac", "anise", "antelope", "arbor", "arena", "armadillo", + "avocet", "azalea", "balsam", "bayou", "beacon", "blizzard", "bluebell", "bluebird", + "bluejay", "bobolink", "borage", "boreal", "buckeye", "buckthorn", "buttercup", + "cabana", "calico", "canopy", "caraway", "cardamom", "cattail", "celadon", "centaur", + "chambray", "chamois", "champlain", "chestnuts", "chickadee", "chinook", "chipmunk", "cinnabar", + "cirrus", "citrine", "clematis", "copperhead", + "crocodile", "currant", "cuttlebone", "daffy", "dapple", "delphinium", "dervish", "diamondback", + "dogwood", "dolphins", "dragonfly", "driftwood", "dusk", "dustpan", "ebony", "edelweiss", + "emperor", "endive", "estuary", "everglade", "fairway", "feldspar", "fennel", "fieldstone", + "firebrand", "firefly", "fireweed", "firework", "flagstone", "fossil", "frostbite", "galleon", + "gardener", "geranium", "gingko", "ginseng", "goldfish", "goldfinch", "goldenrod", "graphite", + "greenfinch", "guppy", "haiku", "halibut", "hammerhead", "harbinger", "harvest", "hatchling", + "havana", "hawthorn", "hazelnut", "heartwood", "henna", "heron", "highrise", "homestead", + "honeycomb", "honeydew", "horseshoe", "hyacinth", "iceland", "icicle", "indigobird", "ironwood", + "jacaranda", "jamboree", "javelina", "jellyfish", "junebug", "kaleido", "kayaker", "kerchief", + "keystone", "kingdom", "labrador", "lacewing", "ladybug", "lakeside", "lamplight", "leopard", + "lighthouse", "lilypad", "lullaby", "magnet", "mahonia", "mandolin", "manzanita", "maraschino", + "mariner", "marsupial", "mastodon", "matterhorn", "mayflower", "mayfly", "meadowlark", "merlot", + "meteor", "midshipman", "millpond", "mimosa", "minnow", "mockingbird", "molten", "monarch", + "monsoon", "moondust", "moonlight", "moorland", "morning", "mossland", "mountain", "mulch", + "narcissus", "nautilus", "nettlebush", "northstar", "nuthatch", "obsidian", "okra", "olivine", + "opalescent", "orchidea", "orchard", "ornament", "outrigger", "oxalis", "paddler", "paintbrush", + "papyrus", "paradise", "pasture", "patchwork", "pathway", "peridot", "periwinkle", "petalbloom", + "petrel", "petunia", "phlox", "pikeperch", "pinecone", "pioneer", "pipevine", "platypus", + "pomelo", "pondweed", "porpoise", "powder", "promise", "puddle", "pumice", "puzzle", + "quetzal", "quicksilver", "racoon", "ragwort", "rainforest", "ramble", "rapid", "rascal", + "raspberry", "redbud", "redfern", "redpoll", "reedling", "ringtail", "riverbed", "riverbird", + "riverstone", "rockcress", "roebuck", "rosebay", "rosehip", "rosemary", "rowan", "rumble", + "runaway", "rustler", "sagebrush", "sailcloth", "salamander", "salsify", "samphire", "sandbar", + "sanddollar", "sandpiper", "santolina", "sapodilla", "sassafras", "scallion", "schooner", "seafoam", + "seafrost", "seagrass", "seahorse", "seaport", "seashell", "seaspray", "shamble", "shimmer", + "shoreline", "silkmoth", "silverfox", "skylark", "snapdragon", "snowberry", "snowdrop", "snowfall", + "snowmelt", "softwood", "songbird", "sorghum", "southwind", "speedwell", "spinnaker", "spruce", + "starlight", "starling", "stormcloud", "summit", "sundance", "sundew", "sundial", "sunflower", + "surface", "swallowtail", "sweetcorn", "sycamore", "tabletop", "tamarack", "tamarind", "tangerine", + "tarragon", "telescope", "thicket", "thrasher", "thunder", "thyme", "tideline", "timberland", + "tinderbox", "topiary", "torchwood", "totem", "tradewind", "treasure", "tremolo", "trinket", + "trumpetvine", "tugboat", "tundra", "turnstone", "underbrush", "vagabond", "valerian", "vanilla", + "velveteen", "vermilion", "vinca", "vineyard", "violet", "voyager", "wagonwheel", "walnutwood", + "watermark", "watershed", "waterway", "wavefront", "westerly", "whaleback", "whetstone", "wicker", + "wildbloom", "wildflower", "wilderness", "windsong", "windward", "winterberry", "woodbine", "woodfern", + "woodland", "woodthrush", "woolgrass", "yellowfin", "zenithal", "zucchini", +} diff --git a/management/server/agentnetwork/manager.go b/management/server/agentnetwork/manager.go new file mode 100644 index 000000000..2a9e4fa87 --- /dev/null +++ b/management/server/agentnetwork/manager.go @@ -0,0 +1,896 @@ +package agentnetwork + +import ( + "context" + "errors" + "fmt" + "math/rand" + "slices" + "strings" + "sync" + "time" + + log "github.com/sirupsen/logrus" + + "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy" + "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey" + "github.com/netbirdio/netbird/management/server/account" + "github.com/netbirdio/netbird/management/server/activity" + "github.com/netbirdio/netbird/management/server/agentnetwork/labelgen" + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/permissions" + "github.com/netbirdio/netbird/management/server/permissions/modules" + "github.com/netbirdio/netbird/management/server/permissions/operations" + "github.com/netbirdio/netbird/management/server/store" + "github.com/netbirdio/netbird/shared/management/proto" + "github.com/netbirdio/netbird/shared/management/status" +) + +// ensureSessionKeys mints an ed25519 session keypair on the provider +// when one is missing. Idempotent: skips when both fields are already +// populated (e.g. update or migrated rows). The keys are used by the +// synthesised reverse-proxy service to sign / verify session JWTs +// after a successful OIDC handshake. +func ensureSessionKeys(p *types.Provider) error { + if p.SessionPrivateKey != "" && p.SessionPublicKey != "" { + return nil + } + pair, err := sessionkey.GenerateKeyPair() + if err != nil { + return fmt.Errorf("generate provider session keys: %w", err) + } + p.SessionPrivateKey = pair.PrivateKey + p.SessionPublicKey = pair.PublicKey + return nil +} + +// Manager governs the lifecycle of Agent Network providers and policies. +type Manager interface { + GetAllProviders(ctx context.Context, accountID, userID string) ([]*types.Provider, error) + GetProvider(ctx context.Context, accountID, userID, providerID string) (*types.Provider, error) + CreateProvider(ctx context.Context, userID string, provider *types.Provider, bootstrapCluster string) (*types.Provider, error) + UpdateProvider(ctx context.Context, userID string, provider *types.Provider) (*types.Provider, error) + DeleteProvider(ctx context.Context, accountID, userID, providerID string) error + + GetAllPolicies(ctx context.Context, accountID, userID string) ([]*types.Policy, error) + GetPolicy(ctx context.Context, accountID, userID, policyID string) (*types.Policy, error) + CreatePolicy(ctx context.Context, userID string, policy *types.Policy) (*types.Policy, error) + UpdatePolicy(ctx context.Context, userID string, policy *types.Policy) (*types.Policy, error) + DeletePolicy(ctx context.Context, accountID, userID, policyID string) error + + GetAllGuardrails(ctx context.Context, accountID, userID string) ([]*types.Guardrail, error) + GetGuardrail(ctx context.Context, accountID, userID, guardrailID string) (*types.Guardrail, error) + CreateGuardrail(ctx context.Context, userID string, guardrail *types.Guardrail) (*types.Guardrail, error) + UpdateGuardrail(ctx context.Context, userID string, guardrail *types.Guardrail) (*types.Guardrail, error) + DeleteGuardrail(ctx context.Context, accountID, userID, guardrailID string) error + + GetAllBudgetRules(ctx context.Context, accountID, userID string) ([]*types.AccountBudgetRule, error) + GetBudgetRule(ctx context.Context, accountID, userID, ruleID string) (*types.AccountBudgetRule, error) + CreateBudgetRule(ctx context.Context, userID string, rule *types.AccountBudgetRule) (*types.AccountBudgetRule, error) + UpdateBudgetRule(ctx context.Context, userID string, rule *types.AccountBudgetRule) (*types.AccountBudgetRule, error) + DeleteBudgetRule(ctx context.Context, accountID, userID, ruleID string) error + + GetSettings(ctx context.Context, accountID, userID string) (*types.Settings, error) + UpdateSettings(ctx context.Context, userID string, settings *types.Settings) (*types.Settings, error) + + ListConsumption(ctx context.Context, accountID, userID string) ([]*types.Consumption, error) + ListAccessLogs(ctx context.Context, accountID, userID string, filter types.AgentNetworkAccessLogFilter) ([]*types.AgentNetworkAccessLog, int64, error) + GetUsageOverview(ctx context.Context, accountID, userID string, filter types.AgentNetworkAccessLogFilter, granularity types.UsageGranularity) ([]*types.AgentNetworkUsageBucket, error) + StartAccessLogCleanup(ctx context.Context, cleanupIntervalHours int) + RecordConsumption(ctx context.Context, accountID string, kind types.ConsumptionDimension, dimID string, windowSeconds, tokensIn, tokensOut int64, costUSD float64) error + RecordAccountBudgetUsage(ctx context.Context, accountID, userID string, groupIDs []string, tokensIn, tokensOut int64, costUSD float64) error + RecordUsage(ctx context.Context, in RecordUsageInput) error + SelectPolicyForRequest(ctx context.Context, in PolicySelectionInput) (*PolicySelectionResult, error) +} + +// PolicySelectionInput is the per-request selection envelope. The +// proxy populates it from CapturedData (account, user, groups) plus +// the provider llm_router resolved. +type PolicySelectionInput struct { + AccountID string + UserID string + GroupIDs []string + ProviderID string +} + +// PolicySelectionResult names the policy that "pays" for this request +// plus the deny envelope when every applicable policy has exhausted +// every cap. AttributionGroupID is the lowest group id (string sort) +// of caller_groups ∩ selected_policy.source_groups; empty when no +// group dimension applies. WindowSeconds is the chosen policy's +// effective window length in seconds (token_limit's wins when both +// halves are enabled with mismatched windows; budget_limit's +// otherwise; 0 when no caps are configured at all). +type PolicySelectionResult struct { + Allow bool + SelectedPolicyID string + AttributionGroupID string + WindowSeconds int64 + DenyCode string + DenyReason string +} + +type managerImpl struct { + store store.Store + accountManager account.Manager + permissionsManager permissions.Manager + proxyController proxy.Controller + + // reconcileCache holds the last set of synthesised proxy mappings + // per account so reconcile can emit precise Create/Update/Delete + // updates instead of a full re-push on every mutation. Keyed by + // accountID, then by synthesised service ID. + reconcileMu sync.Mutex + reconcileCache map[string]map[string]*proto.ProxyMapping + + // labelRngMu guards labelRng. PickUnique consumes math/rand.Source + // state; concurrent provider creates would otherwise race. + labelRngMu sync.Mutex + labelRng *rand.Rand +} + +// NewManager constructs the persistent Agent Network manager. The +// manager persists provider/policy/guardrail configuration and, on +// every mutation, reconciles the in-memory synthesised reverse-proxy +// services with the proxy cluster via proxyController. Pass nil for +// proxyController to disable the reconcile push (useful in tests). +func NewManager( + store store.Store, + permissionsManager permissions.Manager, + accountManager account.Manager, + proxyController proxy.Controller, +) Manager { + return &managerImpl{ + store: store, + accountManager: accountManager, + permissionsManager: permissionsManager, + proxyController: proxyController, + reconcileCache: make(map[string]map[string]*proto.ProxyMapping), + labelRng: rand.New(rand.NewSource(time.Now().UnixNano())), + } +} + +func (m *managerImpl) GetAllProviders(ctx context.Context, accountID, userID string) ([]*types.Provider, error) { + if err := m.requirePermission(ctx, accountID, userID, operations.Read); err != nil { + return nil, err + } + return m.store.GetAccountAgentNetworkProviders(ctx, store.LockingStrengthNone, accountID) +} + +func (m *managerImpl) GetProvider(ctx context.Context, accountID, userID, providerID string) (*types.Provider, error) { + if err := m.requirePermission(ctx, accountID, userID, operations.Read); err != nil { + return nil, err + } + return m.store.GetAgentNetworkProviderByID(ctx, store.LockingStrengthNone, accountID, providerID) +} + +// CreateProvider persists a new provider for the account. bootstrapCluster +// is used only when the per-account agent-network Settings row hasn't +// been created yet; otherwise it is ignored (the cluster is pinned on +// Settings and every provider in the account routes through it). +func (m *managerImpl) CreateProvider(ctx context.Context, userID string, provider *types.Provider, bootstrapCluster string) (*types.Provider, error) { + if err := m.requirePermission(ctx, provider.AccountID, userID, operations.Create); err != nil { + return nil, err + } + + // An empty api_key would silently produce a synthesised service + // that 401s on every upstream request. Surface the misconfiguration + // at create time instead. + if strings.TrimSpace(provider.APIKey) == "" { + return nil, status.Errorf(status.InvalidArgument, "api_key is required when creating an agent network provider") + } + + if provider.ID == "" { + fresh := types.NewProvider(provider.AccountID) + provider.ID = fresh.ID + provider.CreatedAt = fresh.CreatedAt + provider.UpdatedAt = fresh.UpdatedAt + } + + if err := ensureSessionKeys(provider); err != nil { + return nil, err + } + + if err := m.store.SaveAgentNetworkProvider(ctx, provider); err != nil { + return nil, fmt.Errorf("save agent network provider: %w", err) + } + + if strings.TrimSpace(bootstrapCluster) != "" { + if _, err := m.bootstrapSettingsIfNeeded(ctx, provider.AccountID, bootstrapCluster); err != nil { + // The provider create has already succeeded; logging the + // bootstrap miss matches the plan's PoC behaviour. The synth + // path treats a missing settings row as a no-op, and the next + // provider create retries the bootstrap. + log.WithContext(ctx).Debugf("agent-network bootstrap settings for account %s on cluster %s: %v", provider.AccountID, bootstrapCluster, err) + } + } + + m.accountManager.StoreEvent(ctx, userID, provider.ID, provider.AccountID, activity.AgentNetworkProviderCreated, provider.EventMeta()) + m.reconcile(ctx, provider.AccountID) + + return provider, nil +} + +func (m *managerImpl) UpdateProvider(ctx context.Context, userID string, provider *types.Provider) (*types.Provider, error) { + if err := m.requirePermission(ctx, provider.AccountID, userID, operations.Update); err != nil { + return nil, err + } + + existing, err := m.store.GetAgentNetworkProviderByID(ctx, store.LockingStrengthUpdate, provider.AccountID, provider.ID) + if err != nil { + return nil, fmt.Errorf("failed to get agent network provider: %w", err) + } + + // Preserve the API key if the caller didn't rotate it. A + // whitespace-only value is treated as "not rotated" rather than a + // real key, but it must not silently overwrite a valid stored key. + if provider.APIKey == "" { + provider.APIKey = existing.APIKey + } else if strings.TrimSpace(provider.APIKey) == "" { + return nil, status.Errorf(status.InvalidArgument, "api_key must be non-blank when rotating an agent network provider") + } + // Always preserve the session keypair across updates so existing + // session cookies stay valid. The keys are server-managed and + // never surfaced through the API. + provider.SessionPrivateKey = existing.SessionPrivateKey + provider.SessionPublicKey = existing.SessionPublicKey + if err := ensureSessionKeys(provider); err != nil { + return nil, err + } + provider.CreatedAt = existing.CreatedAt + provider.UpdatedAt = time.Now().UTC() + + if err := m.store.SaveAgentNetworkProvider(ctx, provider); err != nil { + return nil, fmt.Errorf("save agent network provider: %w", err) + } + + m.accountManager.StoreEvent(ctx, userID, provider.ID, provider.AccountID, activity.AgentNetworkProviderUpdated, provider.EventMeta()) + m.reconcile(ctx, provider.AccountID) + + return provider, nil +} + +func (m *managerImpl) DeleteProvider(ctx context.Context, accountID, userID, providerID string) error { + if err := m.requirePermission(ctx, accountID, userID, operations.Delete); err != nil { + return err + } + + provider, err := m.store.GetAgentNetworkProviderByID(ctx, store.LockingStrengthUpdate, accountID, providerID) + if err != nil { + return fmt.Errorf("failed to get agent network provider: %w", err) + } + + // Refuse to delete while any policy still references this provider. + // The operator must detach it first. + policies, err := m.store.GetAccountAgentNetworkPolicies(ctx, store.LockingStrengthNone, accountID) + if err != nil { + return fmt.Errorf("failed to get agent network policies: %w", err) + } + var blocking []string + for _, p := range policies { + if slices.Contains(p.DestinationProviderIDs, providerID) { + blocking = append(blocking, p.Name) + } + } + if len(blocking) > 0 { + return status.Errorf( + status.InvalidArgument, + "provider is in use by %d %s (%s); detach it before deleting", + len(blocking), + pluralize(len(blocking), "policy", "policies"), + strings.Join(blocking, ", "), + ) + } + + if err := m.store.DeleteAgentNetworkProvider(ctx, accountID, providerID); err != nil { + return fmt.Errorf("failed to delete agent network provider: %w", err) + } + + m.accountManager.StoreEvent(ctx, userID, providerID, accountID, activity.AgentNetworkProviderDeleted, provider.EventMeta()) + m.reconcile(ctx, accountID) + + return nil +} + +func pluralize(n int, singular, plural string) string { + if n == 1 { + return singular + } + return plural +} + +func (m *managerImpl) GetAllPolicies(ctx context.Context, accountID, userID string) ([]*types.Policy, error) { + if err := m.requirePermission(ctx, accountID, userID, operations.Read); err != nil { + return nil, err + } + return m.store.GetAccountAgentNetworkPolicies(ctx, store.LockingStrengthNone, accountID) +} + +func (m *managerImpl) GetPolicy(ctx context.Context, accountID, userID, policyID string) (*types.Policy, error) { + if err := m.requirePermission(ctx, accountID, userID, operations.Read); err != nil { + return nil, err + } + return m.store.GetAgentNetworkPolicyByID(ctx, store.LockingStrengthNone, accountID, policyID) +} + +func (m *managerImpl) CreatePolicy(ctx context.Context, userID string, policy *types.Policy) (*types.Policy, error) { + if err := m.requirePermission(ctx, policy.AccountID, userID, operations.Create); err != nil { + return nil, err + } + + if policy.ID == "" { + fresh := types.NewPolicy(policy.AccountID) + policy.ID = fresh.ID + policy.CreatedAt = fresh.CreatedAt + policy.UpdatedAt = fresh.UpdatedAt + } + + if err := m.validateProviderRefs(ctx, policy.AccountID, policy.DestinationProviderIDs); err != nil { + return nil, err + } + + if err := m.store.SaveAgentNetworkPolicy(ctx, policy); err != nil { + return nil, fmt.Errorf("failed to save agent network policy: %w", err) + } + + m.accountManager.StoreEvent(ctx, userID, policy.ID, policy.AccountID, activity.AgentNetworkPolicyCreated, policy.EventMeta()) + m.reconcile(ctx, policy.AccountID) + + return policy, nil +} + +func (m *managerImpl) UpdatePolicy(ctx context.Context, userID string, policy *types.Policy) (*types.Policy, error) { + if err := m.requirePermission(ctx, policy.AccountID, userID, operations.Update); err != nil { + return nil, err + } + + existing, err := m.store.GetAgentNetworkPolicyByID(ctx, store.LockingStrengthUpdate, policy.AccountID, policy.ID) + if err != nil { + return nil, fmt.Errorf("failed to get agent network policy: %w", err) + } + + if err := m.validateProviderRefs(ctx, policy.AccountID, policy.DestinationProviderIDs); err != nil { + return nil, err + } + + policy.CreatedAt = existing.CreatedAt + policy.UpdatedAt = time.Now().UTC() + + if err := m.store.SaveAgentNetworkPolicy(ctx, policy); err != nil { + return nil, fmt.Errorf("failed to save agent network policy: %w", err) + } + + m.accountManager.StoreEvent(ctx, userID, policy.ID, policy.AccountID, activity.AgentNetworkPolicyUpdated, policy.EventMeta()) + m.reconcile(ctx, policy.AccountID) + + return policy, nil +} + +func (m *managerImpl) DeletePolicy(ctx context.Context, accountID, userID, policyID string) error { + if err := m.requirePermission(ctx, accountID, userID, operations.Delete); err != nil { + return err + } + + policy, err := m.store.GetAgentNetworkPolicyByID(ctx, store.LockingStrengthUpdate, accountID, policyID) + if err != nil { + return fmt.Errorf("failed to get agent network policy: %w", err) + } + + if err := m.store.DeleteAgentNetworkPolicy(ctx, accountID, policyID); err != nil { + return fmt.Errorf("failed to delete agent network policy: %w", err) + } + + m.accountManager.StoreEvent(ctx, userID, policyID, accountID, activity.AgentNetworkPolicyDeleted, policy.EventMeta()) + m.reconcile(ctx, accountID) + + return nil +} + +func (m *managerImpl) GetAllGuardrails(ctx context.Context, accountID, userID string) ([]*types.Guardrail, error) { + if err := m.requirePermission(ctx, accountID, userID, operations.Read); err != nil { + return nil, err + } + return m.store.GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, accountID) +} + +func (m *managerImpl) GetGuardrail(ctx context.Context, accountID, userID, guardrailID string) (*types.Guardrail, error) { + if err := m.requirePermission(ctx, accountID, userID, operations.Read); err != nil { + return nil, err + } + return m.store.GetAgentNetworkGuardrailByID(ctx, store.LockingStrengthNone, accountID, guardrailID) +} + +func (m *managerImpl) CreateGuardrail(ctx context.Context, userID string, guardrail *types.Guardrail) (*types.Guardrail, error) { + if err := m.requirePermission(ctx, guardrail.AccountID, userID, operations.Create); err != nil { + return nil, err + } + + if guardrail.ID == "" { + fresh := types.NewGuardrail(guardrail.AccountID) + guardrail.ID = fresh.ID + guardrail.CreatedAt = fresh.CreatedAt + guardrail.UpdatedAt = fresh.UpdatedAt + } + + if err := m.store.SaveAgentNetworkGuardrail(ctx, guardrail); err != nil { + return nil, fmt.Errorf("failed to save agent network guardrail: %w", err) + } + + m.accountManager.StoreEvent(ctx, userID, guardrail.ID, guardrail.AccountID, activity.AgentNetworkGuardrailCreated, guardrail.EventMeta()) + m.reconcile(ctx, guardrail.AccountID) + + return guardrail, nil +} + +func (m *managerImpl) UpdateGuardrail(ctx context.Context, userID string, guardrail *types.Guardrail) (*types.Guardrail, error) { + if err := m.requirePermission(ctx, guardrail.AccountID, userID, operations.Update); err != nil { + return nil, err + } + + existing, err := m.store.GetAgentNetworkGuardrailByID(ctx, store.LockingStrengthUpdate, guardrail.AccountID, guardrail.ID) + if err != nil { + return nil, fmt.Errorf("failed to get agent network guardrail: %w", err) + } + + guardrail.CreatedAt = existing.CreatedAt + guardrail.UpdatedAt = time.Now().UTC() + + if err := m.store.SaveAgentNetworkGuardrail(ctx, guardrail); err != nil { + return nil, fmt.Errorf("failed to save agent network guardrail: %w", err) + } + + m.accountManager.StoreEvent(ctx, userID, guardrail.ID, guardrail.AccountID, activity.AgentNetworkGuardrailUpdated, guardrail.EventMeta()) + m.reconcile(ctx, guardrail.AccountID) + + return guardrail, nil +} + +func (m *managerImpl) DeleteGuardrail(ctx context.Context, accountID, userID, guardrailID string) error { + if err := m.requirePermission(ctx, accountID, userID, operations.Delete); err != nil { + return err + } + + guardrail, err := m.store.GetAgentNetworkGuardrailByID(ctx, store.LockingStrengthUpdate, accountID, guardrailID) + if err != nil { + return fmt.Errorf("failed to get agent network guardrail: %w", err) + } + + if err := m.store.DeleteAgentNetworkGuardrail(ctx, accountID, guardrailID); err != nil { + return fmt.Errorf("failed to delete agent network guardrail: %w", err) + } + + m.accountManager.StoreEvent(ctx, userID, guardrailID, accountID, activity.AgentNetworkGuardrailDeleted, guardrail.EventMeta()) + m.reconcile(ctx, accountID) + + return nil +} + +// GetAllBudgetRules returns every account-level budget rule for the account. +func (m *managerImpl) GetAllBudgetRules(ctx context.Context, accountID, userID string) ([]*types.AccountBudgetRule, error) { + if err := m.requirePermission(ctx, accountID, userID, operations.Read); err != nil { + return nil, err + } + return m.store.GetAccountAgentNetworkBudgetRules(ctx, store.LockingStrengthNone, accountID) +} + +// GetBudgetRule returns a single account-level budget rule. +func (m *managerImpl) GetBudgetRule(ctx context.Context, accountID, userID, ruleID string) (*types.AccountBudgetRule, error) { + if err := m.requirePermission(ctx, accountID, userID, operations.Read); err != nil { + return nil, err + } + return m.store.GetAgentNetworkBudgetRuleByID(ctx, store.LockingStrengthNone, accountID, ruleID) +} + +// CreateBudgetRule persists a new account-level budget rule. Budget rules are +// enforced at request time (CheckLLMPolicyLimits), not baked into the synth +// proxy config, so no reconcile is needed. +func (m *managerImpl) CreateBudgetRule(ctx context.Context, userID string, rule *types.AccountBudgetRule) (*types.AccountBudgetRule, error) { + if err := m.requirePermission(ctx, rule.AccountID, userID, operations.Create); err != nil { + return nil, err + } + + if rule.ID == "" { + fresh := types.NewAccountBudgetRule(rule.AccountID) + rule.ID = fresh.ID + rule.CreatedAt = fresh.CreatedAt + rule.UpdatedAt = fresh.UpdatedAt + } + + if err := m.store.SaveAgentNetworkBudgetRule(ctx, rule); err != nil { + return nil, fmt.Errorf("save agent network budget rule: %w", err) + } + + m.accountManager.StoreEvent(ctx, userID, rule.ID, rule.AccountID, activity.AgentNetworkBudgetRuleCreated, rule.EventMeta()) + + return rule, nil +} + +// UpdateBudgetRule updates an existing account-level budget rule. +func (m *managerImpl) UpdateBudgetRule(ctx context.Context, userID string, rule *types.AccountBudgetRule) (*types.AccountBudgetRule, error) { + if err := m.requirePermission(ctx, rule.AccountID, userID, operations.Update); err != nil { + return nil, err + } + + existing, err := m.store.GetAgentNetworkBudgetRuleByID(ctx, store.LockingStrengthUpdate, rule.AccountID, rule.ID) + if err != nil { + return nil, fmt.Errorf("get agent network budget rule: %w", err) + } + + rule.CreatedAt = existing.CreatedAt + rule.UpdatedAt = time.Now().UTC() + + if err := m.store.SaveAgentNetworkBudgetRule(ctx, rule); err != nil { + return nil, fmt.Errorf("save agent network budget rule: %w", err) + } + + m.accountManager.StoreEvent(ctx, userID, rule.ID, rule.AccountID, activity.AgentNetworkBudgetRuleUpdated, rule.EventMeta()) + + return rule, nil +} + +// DeleteBudgetRule removes an account-level budget rule. +func (m *managerImpl) DeleteBudgetRule(ctx context.Context, accountID, userID, ruleID string) error { + if err := m.requirePermission(ctx, accountID, userID, operations.Delete); err != nil { + return err + } + + rule, err := m.store.GetAgentNetworkBudgetRuleByID(ctx, store.LockingStrengthUpdate, accountID, ruleID) + if err != nil { + return fmt.Errorf("get agent network budget rule: %w", err) + } + + if err := m.store.DeleteAgentNetworkBudgetRule(ctx, accountID, ruleID); err != nil { + return fmt.Errorf("delete agent network budget rule: %w", err) + } + + m.accountManager.StoreEvent(ctx, userID, ruleID, accountID, activity.AgentNetworkBudgetRuleDeleted, rule.EventMeta()) + + return nil +} + +// UpdateSettings applies the mutable account-level settings — the collection +// toggles — onto the existing row. Cluster and Subdomain are immutable and are +// preserved from the persisted row regardless of the input. Because the +// collection toggles change the synthesised service config (prompt-capture +// gating, access-log emission), a reconcile is triggered so the proxy and peer +// network maps converge on the new state. +func (m *managerImpl) UpdateSettings(ctx context.Context, userID string, settings *types.Settings) (*types.Settings, error) { + if err := m.requirePermission(ctx, settings.AccountID, userID, operations.Update); err != nil { + return nil, err + } + + existing, err := m.store.GetAgentNetworkSettings(ctx, store.LockingStrengthUpdate, settings.AccountID) + if err != nil { + return nil, fmt.Errorf("get agent network settings: %w", err) + } + + existing.EnableLogCollection = settings.EnableLogCollection + existing.EnablePromptCollection = settings.EnablePromptCollection + existing.RedactPii = settings.RedactPii + existing.AccessLogRetentionDays = settings.AccessLogRetentionDays + existing.UpdatedAt = time.Now().UTC() + + if err := m.store.SaveAgentNetworkSettings(ctx, existing); err != nil { + return nil, fmt.Errorf("save agent network settings: %w", err) + } + + m.accountManager.StoreEvent(ctx, userID, settings.AccountID, settings.AccountID, activity.AgentNetworkSettingsUpdated, map[string]any{ + "log_collection": existing.EnableLogCollection, + "prompt_collection": existing.EnablePromptCollection, + "redact_pii": existing.RedactPii, + }) + m.reconcile(ctx, settings.AccountID) + + return existing, nil +} + +// validateProviderRefs ensures every destination provider id refers to a +// provider that exists in the same account. +func (m *managerImpl) validateProviderRefs(ctx context.Context, accountID string, providerIDs []string) error { + if len(providerIDs) == 0 { + return nil + } + for _, id := range providerIDs { + if _, err := m.store.GetAgentNetworkProviderByID(ctx, store.LockingStrengthNone, accountID, id); err != nil { + // Only a genuine not-found means the reference is invalid; a + // store/runtime error must propagate as-is rather than be + // masked as a client validation error. + var sErr *status.Error + if errors.As(err, &sErr) && sErr.Type() == status.NotFound { + return status.Errorf(status.InvalidArgument, "destination_provider_ids: provider %s does not exist", id) + } + return fmt.Errorf("get destination provider %s: %w", id, err) + } + } + return nil +} + +// GetSettings returns the agent-network settings row for the account. +// Returns the underlying status.NotFound when no row has been +// bootstrapped yet (i.e. the account has no providers). +func (m *managerImpl) GetSettings(ctx context.Context, accountID, userID string) (*types.Settings, error) { + if err := m.requirePermission(ctx, accountID, userID, operations.Read); err != nil { + return nil, err + } + return m.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID) +} + +// bootstrapSettingsIfNeeded creates the per-account agent-network +// settings row when missing. The cluster comes from the create-time +// hint the dashboard sends (auto-picked from the active cluster list); +// the subdomain is picked from the curated wordlist avoiding +// collisions on the same cluster. Idempotent: if a row already exists +// it is returned untouched and the hint is ignored. +func (m *managerImpl) bootstrapSettingsIfNeeded(ctx context.Context, accountID, providerCluster string) (*types.Settings, error) { + if accountID == "" { + return nil, fmt.Errorf("bootstrap settings: account id is required") + } + if strings.TrimSpace(providerCluster) == "" { + return nil, fmt.Errorf("bootstrap settings: provider cluster is required") + } + + existing, err := m.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID) + if err == nil { + return existing, nil + } + var sErr *status.Error + if !errors.As(err, &sErr) || sErr.Type() != status.NotFound { + return nil, fmt.Errorf("get agent network settings: %w", err) + } + + siblings, err := m.store.GetAgentNetworkSettingsByCluster(ctx, store.LockingStrengthNone, providerCluster) + if err != nil { + return nil, fmt.Errorf("list agent network settings on cluster: %w", err) + } + taken := make(map[string]struct{}, len(siblings)) + for _, s := range siblings { + taken[s.Subdomain] = struct{}{} + } + + suffix := accountID + if len(suffix) > 4 { + suffix = suffix[:4] + } + + m.labelRngMu.Lock() + subdomain := labelgen.PickUnique(m.labelRng, taken, suffix) + m.labelRngMu.Unlock() + + now := time.Now().UTC() + settings := &types.Settings{ + AccountID: accountID, + Cluster: providerCluster, + Subdomain: subdomain, + // Logs on by default; usage is collected regardless. Retention bounds + // how long full log rows are kept. + EnableLogCollection: true, + AccessLogRetentionDays: types.DefaultAccessLogRetentionDays, + CreatedAt: now, + UpdatedAt: now, + } + if err := m.store.SaveAgentNetworkSettings(ctx, settings); err != nil { + return nil, fmt.Errorf("save agent network settings: %w", err) + } + return settings, nil +} + +// ListConsumption returns every consumption row recorded for the +// account, ordered window-newest-first. Backs the dashboard's basic +// counter view; permission gate is the same Read role that gates +// every other agent-network surface. +func (m *managerImpl) ListConsumption(ctx context.Context, accountID, userID string) ([]*types.Consumption, error) { + if err := m.requirePermission(ctx, accountID, userID, operations.Read); err != nil { + return nil, err + } + return m.store.ListAgentNetworkConsumption(ctx, store.LockingStrengthNone, accountID) +} + +// ListAccessLogs returns a paginated, server-side-filtered page of +// agent-network access logs plus the total count matching the filter. +func (m *managerImpl) ListAccessLogs(ctx context.Context, accountID, userID string, filter types.AgentNetworkAccessLogFilter) ([]*types.AgentNetworkAccessLog, int64, error) { + if err := m.requirePermission(ctx, accountID, userID, operations.Read); err != nil { + return nil, 0, err + } + return m.store.GetAgentNetworkAccessLogs(ctx, store.LockingStrengthNone, accountID, filter) +} + +// GetUsageOverview returns the filtered usage rows aggregated into time buckets +// at the requested granularity, oldest-first. +func (m *managerImpl) GetUsageOverview(ctx context.Context, accountID, userID string, filter types.AgentNetworkAccessLogFilter, granularity types.UsageGranularity) ([]*types.AgentNetworkUsageBucket, error) { + if err := m.requirePermission(ctx, accountID, userID, operations.Read); err != nil { + return nil, err + } + rows, err := m.store.GetAgentNetworkUsageRows(ctx, store.LockingStrengthNone, accountID, filter) + if err != nil { + return nil, err + } + return types.AggregateUsageByGranularity(rows, granularity), nil +} + +// StartAccessLogCleanup launches a background sweep that periodically deletes +// each account's agent-network access-log rows older than that account's +// AccessLogRetentionDays. Usage records are never swept. A non-positive +// interval defaults to 24h. +func (m *managerImpl) StartAccessLogCleanup(ctx context.Context, cleanupIntervalHours int) { + if cleanupIntervalHours <= 0 { + cleanupIntervalHours = 24 + } + interval := time.Duration(cleanupIntervalHours) * time.Hour + + go func() { + ticker := time.NewTicker(interval) + defer ticker.Stop() + + m.cleanupAccessLogsOnce(ctx) // run once on startup + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + m.cleanupAccessLogsOnce(ctx) + } + } + }() +} + +// cleanupAccessLogsOnce sweeps every account's expired access-log rows against +// its configured retention. Best-effort: a per-account failure is logged and +// the sweep continues. +func (m *managerImpl) cleanupAccessLogsOnce(ctx context.Context) { + settings, err := m.store.GetAllAgentNetworkSettings(ctx, store.LockingStrengthNone) + if err != nil { + log.WithContext(ctx).Errorf("agent-network access-log cleanup: list settings: %v", err) + return + } + for _, s := range settings { + if s.AccessLogRetentionDays <= 0 { + continue // keep indefinitely + } + cutoff := time.Now().UTC().AddDate(0, 0, -s.AccessLogRetentionDays) + deleted, err := m.store.DeleteOldAgentNetworkAccessLogs(ctx, s.AccountID, cutoff) + if err != nil { + log.WithContext(ctx).Warnf("agent-network access-log cleanup for account %s: %v", s.AccountID, err) + continue + } + if deleted > 0 { + log.WithContext(ctx).Infof("agent-network access-log cleanup: deleted %d rows for account %s (retention %d days)", deleted, s.AccountID, s.AccessLogRetentionDays) + } + } +} + +// RecordConsumption increments the (dim, window) counter by the +// supplied deltas. The window_start is computed from time.Now under +// the supplied window_seconds so callers don't have to pre-align — +// the proxy's post-flight path simply hands us tokens + cost and +// which dimension we're booking against. +func (m *managerImpl) RecordConsumption(ctx context.Context, accountID string, kind types.ConsumptionDimension, dimID string, windowSeconds, 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") + } + windowStart := types.WindowStart(time.Now(), windowSeconds) + return m.store.IncrementAgentNetworkConsumption(ctx, accountID, kind, dimID, windowSeconds, windowStart, tokensIn, tokensOut, costUSD) +} + +func (m *managerImpl) requirePermission(ctx context.Context, accountID, userID string, op operations.Operation) error { + ok, _, err := m.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.AgentNetwork, op) + if err != nil { + return status.NewPermissionValidationError(err) + } + if !ok { + return status.NewPermissionDeniedError() + } + return nil +} + +type mockManager struct{} + +// NewManagerMock returns a no-op manager useful for tests. +func NewManagerMock() Manager { + return &mockManager{} +} + +func (*mockManager) GetAllProviders(_ context.Context, _, _ string) ([]*types.Provider, error) { + return []*types.Provider{}, nil +} + +func (*mockManager) GetProvider(_ context.Context, _, _, _ string) (*types.Provider, error) { + return &types.Provider{}, nil +} + +func (*mockManager) CreateProvider(_ context.Context, _ string, p *types.Provider, _ string) (*types.Provider, error) { + return p, nil +} + +func (*mockManager) UpdateProvider(_ context.Context, _ string, p *types.Provider) (*types.Provider, error) { + return p, nil +} + +func (*mockManager) DeleteProvider(_ context.Context, _, _, _ string) error { return nil } + +func (*mockManager) GetAllPolicies(_ context.Context, _, _ string) ([]*types.Policy, error) { + return []*types.Policy{}, nil +} + +func (*mockManager) GetPolicy(_ context.Context, _, _, _ string) (*types.Policy, error) { + return &types.Policy{}, nil +} + +func (*mockManager) CreatePolicy(_ context.Context, _ string, p *types.Policy) (*types.Policy, error) { + return p, nil +} + +func (*mockManager) UpdatePolicy(_ context.Context, _ string, p *types.Policy) (*types.Policy, error) { + return p, nil +} + +func (*mockManager) DeletePolicy(_ context.Context, _, _, _ string) error { return nil } + +func (*mockManager) GetAllGuardrails(_ context.Context, _, _ string) ([]*types.Guardrail, error) { + return []*types.Guardrail{}, nil +} + +func (*mockManager) GetGuardrail(_ context.Context, _, _, _ string) (*types.Guardrail, error) { + return &types.Guardrail{}, nil +} + +func (*mockManager) CreateGuardrail(_ context.Context, _ string, g *types.Guardrail) (*types.Guardrail, error) { + return g, nil +} + +func (*mockManager) UpdateGuardrail(_ context.Context, _ string, g *types.Guardrail) (*types.Guardrail, error) { + return g, nil +} + +func (*mockManager) DeleteGuardrail(_ context.Context, _, _, _ string) error { return nil } + +func (*mockManager) GetAllBudgetRules(_ context.Context, _, _ string) ([]*types.AccountBudgetRule, error) { + return []*types.AccountBudgetRule{}, nil +} + +func (*mockManager) GetBudgetRule(_ context.Context, _, _, _ string) (*types.AccountBudgetRule, error) { + return &types.AccountBudgetRule{}, nil +} + +func (*mockManager) CreateBudgetRule(_ context.Context, _ string, r *types.AccountBudgetRule) (*types.AccountBudgetRule, error) { + return r, nil +} + +func (*mockManager) UpdateBudgetRule(_ context.Context, _ string, r *types.AccountBudgetRule) (*types.AccountBudgetRule, error) { + return r, nil +} + +func (*mockManager) DeleteBudgetRule(_ context.Context, _, _, _ string) error { return nil } + +func (*mockManager) GetSettings(_ context.Context, _, _ string) (*types.Settings, error) { + return nil, status.Errorf(status.NotFound, "agent network settings not found") +} + +func (*mockManager) UpdateSettings(_ context.Context, _ string, s *types.Settings) (*types.Settings, error) { + return s, nil +} + +func (*mockManager) ListConsumption(_ context.Context, _, _ string) ([]*types.Consumption, error) { + return nil, nil +} + +func (*mockManager) ListAccessLogs(_ context.Context, _, _ string, _ types.AgentNetworkAccessLogFilter) ([]*types.AgentNetworkAccessLog, int64, error) { + return nil, 0, nil +} + +func (*mockManager) GetUsageOverview(_ context.Context, _, _ string, _ types.AgentNetworkAccessLogFilter, _ types.UsageGranularity) ([]*types.AgentNetworkUsageBucket, error) { + return nil, nil +} + +func (*mockManager) StartAccessLogCleanup(_ context.Context, _ int) {} + +func (*mockManager) RecordConsumption(_ context.Context, _ string, _ types.ConsumptionDimension, _ string, _, _, _ int64, _ float64) error { + return nil +} + +func (*mockManager) RecordAccountBudgetUsage(_ context.Context, _, _ string, _ []string, _, _ int64, _ float64) error { + return nil +} + +func (*mockManager) RecordUsage(_ context.Context, _ RecordUsageInput) error { + return nil +} diff --git a/management/server/agentnetwork/policyselect.go b/management/server/agentnetwork/policyselect.go new file mode 100644 index 000000000..baec154e2 --- /dev/null +++ b/management/server/agentnetwork/policyselect.go @@ -0,0 +1,660 @@ +package agentnetwork + +import ( + "context" + "fmt" + "math" + "sort" + "time" + + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/store" + "github.com/netbirdio/netbird/shared/management/status" +) + +// validateUsageDeltas rejects negative or non-finite usage counters before they +// reach the consumption store, so a bad delta can't decrement or poison totals. +// The store batch method enforces the same invariant; this is the manager-level +// guard so direct callers fail fast with a clear error. +func validateUsageDeltas(tokensIn, tokensOut int64, costUSD float64) error { + if tokensIn < 0 || tokensOut < 0 || costUSD < 0 || math.IsNaN(costUSD) || math.IsInf(costUSD, 0) { + return status.Errorf(status.InvalidArgument, "usage deltas must be non-negative and finite") + } + return nil +} + +// Deny codes the proxy surfaces back to the caller when every +// applicable policy is exhausted. The proxy converts these into +// upstream-shaped error responses. +const ( + //nolint:gosec // policy deny code label, not a credential + denyCodeTokenCapExceeded = "llm_policy.token_cap_exceeded" + //nolint:gosec // policy deny code label, not a credential + denyCodeBudgetCapExceeded = "llm_policy.budget_cap_exceeded" + //nolint:gosec // account deny code label, not a credential + denyCodeAccountTokenCapExceeded = "llm_account.token_cap_exceeded" + //nolint:gosec // account deny code label, not a credential + denyCodeAccountBudgetCapExceeded = "llm_account.budget_cap_exceeded" +) + +// consumptionCache holds the consumption counters prefetched for one +// policy-selection request, keyed by ConsumptionKey. A miss returns a zero +// counter — the same contract the store's single-row getter uses for absent +// rows — so the eval logic is identical whether a counter exists yet or not. +type consumptionCache map[types.ConsumptionKey]*types.Consumption + +func (c consumptionCache) get(accountID string, kind types.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time) *types.Consumption { + key := types.ConsumptionKey{Kind: kind, DimID: dimID, WindowSeconds: windowSeconds, WindowStartUTC: windowStart.UTC()} + if row, ok := c[key]; ok && row != nil { + return row + } + return &types.Consumption{ + AccountID: accountID, + DimensionKind: kind, + DimensionID: dimID, + WindowSeconds: windowSeconds, + WindowStartUTC: windowStart.UTC(), + } +} + +// addLimitKeys records the user/group consumption keys a single enabled (token +// or budget) limit window reads for the given attribution group, into a dedup +// set. attrGroup may be empty (no group dimension applies). +func addLimitKeys(set map[types.ConsumptionKey]struct{}, userID, attrGroup string, windowSeconds int64, now time.Time) { + if windowSeconds <= 0 { + return + } + ws := types.WindowStart(now, windowSeconds) + if userID != "" { + set[types.ConsumptionKey{Kind: types.DimensionUser, DimID: userID, WindowSeconds: windowSeconds, WindowStartUTC: ws}] = struct{}{} + } + if attrGroup != "" { + set[types.ConsumptionKey{Kind: types.DimensionGroup, DimID: attrGroup, WindowSeconds: windowSeconds, WindowStartUTC: ws}] = struct{}{} + } +} + +// prefetchConsumption loads, in one store round-trip, every consumption counter +// that the account-budget ceiling and the candidate policies will read while +// scoring this request. This replaces the per-cap point reads the selector +// previously issued one at a time (the N+1 on the hot path). +func (m *managerImpl) prefetchConsumption(ctx context.Context, in PolicySelectionInput, rules []*types.AccountBudgetRule, candidates []*types.Policy, now time.Time) (consumptionCache, error) { + set := make(map[types.ConsumptionKey]struct{}) + for _, p := range candidates { + attr := lowestIntersect(p.SourceGroups, in.GroupIDs) + if p.Limits.TokenLimit.Enabled { + addLimitKeys(set, in.UserID, attr, p.Limits.TokenLimit.WindowSeconds, now) + } + if p.Limits.BudgetLimit.Enabled { + addLimitKeys(set, in.UserID, attr, p.Limits.BudgetLimit.WindowSeconds, now) + } + } + for _, r := range rules { + if r == nil || !r.Enabled || !budgetRuleApplies(r, in) { + continue + } + attr := lowestIntersect(r.TargetGroups, in.GroupIDs) + if r.Limits.TokenLimit.Enabled { + addLimitKeys(set, in.UserID, attr, r.Limits.TokenLimit.WindowSeconds, now) + } + if r.Limits.BudgetLimit.Enabled { + addLimitKeys(set, in.UserID, attr, r.Limits.BudgetLimit.WindowSeconds, now) + } + } + if len(set) == 0 { + return consumptionCache{}, nil + } + keys := make([]types.ConsumptionKey, 0, len(set)) + for k := range set { + keys = append(keys, k) + } + rows, err := m.store.GetAgentNetworkConsumptionBatch(ctx, store.LockingStrengthNone, in.AccountID, keys) + if err != nil { + return nil, fmt.Errorf("batch read consumption: %w", err) + } + return consumptionCache(rows), nil +} + +// SelectPolicyForRequest picks the policy that "pays" for the +// incoming request. The chosen policy is the one with the largest +// pool that still has headroom — drain the bigger bucket first, +// fall through to the next-biggest only when the current one's +// group cap or shared per-user cap is exhausted. This matches +// operator intuition for layered tiers ("privileged group has the +// 10k budget, regular group has 1k as the safety net") and avoids +// the load-balancer flapping that fraction-based scoring produces +// once any cap has been touched. +// +// Ordering across non-exhausted candidates: +// 1. Policies with NO enabled caps (catch-all-allow) win over any +// capped policy — operators who configure unlimited access +// expect requests to attribute there until they explicitly add +// caps. +// 2. Larger group token cap wins. +// 3. Larger group budget USD cap wins. +// 4. Larger user token cap wins. +// 5. Larger user budget USD cap wins. +// 6. Older created_at wins (deterministic final tiebreak so +// multi-node selection converges). +// +// Returns Allow=true with empty SelectedPolicyID when no policy in +// the account targets the (provider, caller-groups) combination — +// llm_router is the gate that owns "no policy authorises this +// request" semantics; this function trusts that authorisation has +// already happened upstream and only does the limit-aware +// attribution. +func (m *managerImpl) SelectPolicyForRequest(ctx context.Context, in PolicySelectionInput) (*PolicySelectionResult, error) { + if in.AccountID == "" { + return nil, status.Errorf(status.InvalidArgument, "account_id is required") + } + + now := time.Now().UTC() + + rules, err := m.store.GetAccountAgentNetworkBudgetRules(ctx, store.LockingStrengthNone, in.AccountID) + if err != nil { + return nil, fmt.Errorf("list account budget rules: %w", err) + } + policies, err := m.store.GetAccountAgentNetworkPolicies(ctx, store.LockingStrengthNone, in.AccountID) + if err != nil { + return nil, fmt.Errorf("list account policies: %w", err) + } + candidates := filterApplicablePolicies(policies, in) + + // Prefetch every consumption counter the ceiling + candidate policies will + // read, in a single store round-trip, then score against the cache. + cache, err := m.prefetchConsumption(ctx, in, rules, candidates, now) + if err != nil { + return nil, err + } + + // Account-level budget rules are an always-on ceiling, evaluated + // independently of policy selection (they bind even for catch-all-allow + // policies or requests that match no policy). All applicable rules must + // pass — this is where min-wins lives. + if deny, code, reason := checkAccountBudget(in, rules, cache, now); deny { + return &PolicySelectionResult{Allow: false, DenyCode: code, DenyReason: reason}, nil + } + + if len(candidates) == 0 { + return &PolicySelectionResult{Allow: true}, nil + } + scored, lastDenyCode, lastDenyReason := scoreCandidates(in, candidates, cache, now) + if len(scored) == 0 { + return &PolicySelectionResult{ + Allow: false, + DenyCode: lastDenyCode, + DenyReason: lastDenyReason, + }, nil + } + + sort.SliceStable(scored, func(i, j int) bool { + // Catch-all-allow (no caps configured) wins outright over + // any capped policy. + iNoCap := isUncapped(scored[i].policy) + jNoCap := isUncapped(scored[j].policy) + if iNoCap != jNoCap { + return iNoCap + } + // Bigger pool drains first. Group caps dominate (shared + // across the group) before individual caps. + if a, b := groupCapTokens(scored[i].policy), groupCapTokens(scored[j].policy); a != b { + return a > b + } + if a, b := groupCapBudgetUsd(scored[i].policy), groupCapBudgetUsd(scored[j].policy); a != b { + return a > b + } + if a, b := userCapTokens(scored[i].policy), userCapTokens(scored[j].policy); a != b { + return a > b + } + if a, b := userCapBudgetUsd(scored[i].policy), userCapBudgetUsd(scored[j].policy); a != b { + return a > b + } + return scored[i].policy.CreatedAt.Before(scored[j].policy.CreatedAt) + }) + + winner := scored[0] + return &PolicySelectionResult{ + Allow: true, + SelectedPolicyID: winner.policy.ID, + AttributionGroupID: winner.attributionGroup, + WindowSeconds: winner.windowSeconds, + }, nil +} + +// filterApplicablePolicies returns the enabled policies that target +// the requested provider and have at least one of the caller's groups +// in their source_groups. Caller's group set is matched +// case-sensitively against policy.SourceGroups. +func filterApplicablePolicies(policies []*types.Policy, in PolicySelectionInput) []*types.Policy { + if len(policies) == 0 { + return nil + } + groupSet := make(map[string]struct{}, len(in.GroupIDs)) + for _, g := range in.GroupIDs { + if g != "" { + groupSet[g] = struct{}{} + } + } + out := make([]*types.Policy, 0, len(policies)) + for _, p := range policies { + if p == nil || !p.Enabled { + continue + } + if !sliceContains(p.DestinationProviderIDs, in.ProviderID) { + continue + } + if !anyGroupMatches(p.SourceGroups, groupSet) { + continue + } + out = append(out, p) + } + return out +} + +// candidate is the per-policy intermediate the selector ranks. A +// policy that's been exhausted on any enabled cap never makes it +// into this slice; the selector's deny envelope carries the latest +// exhaustion's reason out separately. +type candidate struct { + policy *types.Policy + attributionGroup string + windowSeconds int64 +} + +// scoreCandidates evaluates every applicable policy against the +// caller's current consumption. Exhausted policies are filtered out +// of the returned slice; the most recent exhaustion's deny code + +// human reason is returned alongside so the caller can surface it +// when no candidate survives. +func scoreCandidates( + in PolicySelectionInput, + candidates []*types.Policy, + cache consumptionCache, + now time.Time, +) ([]candidate, string, string) { + out := make([]candidate, 0, len(candidates)) + var lastDenyCode, lastDenyReason string + + for _, p := range candidates { + c, exhausted, denyCode, denyReason := scoreOne(in, p, cache, now) + if exhausted { + lastDenyCode = denyCode + lastDenyReason = denyReason + continue + } + out = append(out, c) + } + return out, lastDenyCode, lastDenyReason +} + +// scoreOne checks a single policy for cap exhaustion. Returns the +// candidate envelope when the policy still has headroom on every +// enabled cap; reports exhausted=true with a deny code naming the +// offending cap kind otherwise. +func scoreOne( + in PolicySelectionInput, + p *types.Policy, + cache consumptionCache, + now time.Time, +) (candidate, bool, string, string) { + attrGroup := lowestIntersect(p.SourceGroups, in.GroupIDs) + c := candidate{ + policy: p, + attributionGroup: attrGroup, + windowSeconds: effectiveWindowSeconds(p), + } + + if p.Limits.TokenLimit.Enabled && p.Limits.TokenLimit.WindowSeconds > 0 { + if exhausted, reason := evalTokenCap(cache, in.AccountID, in.UserID, attrGroup, p.Limits.TokenLimit, now, "policy "+p.ID); exhausted { + return candidate{}, true, denyCodeTokenCapExceeded, reason + } + } + + if p.Limits.BudgetLimit.Enabled && p.Limits.BudgetLimit.WindowSeconds > 0 { + if exhausted, reason := evalBudgetCap(cache, in.AccountID, in.UserID, attrGroup, p.Limits.BudgetLimit, now, "policy "+p.ID); exhausted { + return candidate{}, true, denyCodeBudgetCapExceeded, reason + } + } + + return c, false, "", "" +} + +// evalTokenCap reports whether the token limit is already exhausted for the +// caller in its own window. attrGroup may be empty (no group dimension applies). +// label identifies the cap source ("policy " or "account rule ") for the +// deny reason. It is the shared primitive behind both policy and account-rule +// enforcement. +func evalTokenCap( + cache consumptionCache, + accountID, userID, attrGroup string, + tl types.PolicyTokenLimit, + now time.Time, + label string, +) (bool, string) { + windowStart := types.WindowStart(now, tl.WindowSeconds) + + if tl.UserCap > 0 && userID != "" { + row := cache.get(accountID, types.DimensionUser, userID, tl.WindowSeconds, windowStart) + used := row.TokensInput + row.TokensOutput + if used >= tl.UserCap { + return true, fmt.Sprintf("user token cap exhausted on %s (used %d of %d)", label, used, tl.UserCap) + } + } + + if tl.GroupCap > 0 && attrGroup != "" { + row := cache.get(accountID, types.DimensionGroup, attrGroup, tl.WindowSeconds, windowStart) + used := row.TokensInput + row.TokensOutput + if used >= tl.GroupCap { + return true, fmt.Sprintf("group token cap exhausted on %s (used %d of %d)", label, used, tl.GroupCap) + } + } + + return false, "" +} + +// evalBudgetCap is the budget (USD) counterpart of evalTokenCap. +func evalBudgetCap( + cache consumptionCache, + accountID, userID, attrGroup string, + bl types.PolicyBudgetLimit, + now time.Time, + label string, +) (bool, string) { + windowStart := types.WindowStart(now, bl.WindowSeconds) + + if bl.UserCapUsd > 0 && userID != "" { + row := cache.get(accountID, types.DimensionUser, userID, bl.WindowSeconds, windowStart) + if row.CostUSD >= bl.UserCapUsd { + return true, fmt.Sprintf("user budget cap exhausted on %s (used $%.4f of $%.4f)", label, row.CostUSD, bl.UserCapUsd) + } + } + + if bl.GroupCapUsd > 0 && attrGroup != "" { + row := cache.get(accountID, types.DimensionGroup, attrGroup, bl.WindowSeconds, windowStart) + if row.CostUSD >= bl.GroupCapUsd { + return true, fmt.Sprintf("group budget cap exhausted on %s (used $%.4f of $%.4f)", label, row.CostUSD, bl.GroupCapUsd) + } + } + + return false, "" +} + +// checkAccountBudget evaluates every applicable account-level budget rule as an +// all-must-pass ceiling. A rule applies when the caller is in its TargetUsers, +// one of its TargetGroups, or it has no targets at all (account-wide). Returns +// deny=true with an llm_account.* code on the first exhausted rule. Group caps +// attribute to the lowest intersecting group (the same model policies use), so +// multi-group behavior is unchanged. +func checkAccountBudget(in PolicySelectionInput, rules []*types.AccountBudgetRule, cache consumptionCache, now time.Time) (bool, string, string) { + for _, r := range rules { + if r == nil || !r.Enabled || !budgetRuleApplies(r, in) { + continue + } + attrGroup := lowestIntersect(r.TargetGroups, in.GroupIDs) + label := "account rule " + r.ID + + if r.Limits.TokenLimit.Enabled && r.Limits.TokenLimit.WindowSeconds > 0 { + if exhausted, reason := evalTokenCap(cache, in.AccountID, in.UserID, attrGroup, r.Limits.TokenLimit, now, label); exhausted { + return true, denyCodeAccountTokenCapExceeded, reason + } + } + + if r.Limits.BudgetLimit.Enabled && r.Limits.BudgetLimit.WindowSeconds > 0 { + if exhausted, reason := evalBudgetCap(cache, in.AccountID, in.UserID, attrGroup, r.Limits.BudgetLimit, now, label); exhausted { + return true, denyCodeAccountBudgetCapExceeded, reason + } + } + } + + return false, "", "" +} + +// budgetRuleApplies reports whether an account budget rule binds the caller: +// a direct user match, a group intersection, or an untargeted (account-wide) +// rule. +func budgetRuleApplies(r *types.AccountBudgetRule, in PolicySelectionInput) bool { + if len(r.TargetUsers) == 0 && len(r.TargetGroups) == 0 { + return true + } + if in.UserID != "" && sliceContains(r.TargetUsers, in.UserID) { + return true + } + groupSet := make(map[string]struct{}, len(in.GroupIDs)) + for _, g := range in.GroupIDs { + if g != "" { + groupSet[g] = struct{}{} + } + } + return anyGroupMatches(r.TargetGroups, groupSet) +} + +// RecordAccountBudgetUsage fans the served request's usage out to every +// applicable account budget rule's own (dimension, window) counter. The user +// dimension is always booked when a rule has a user-applicable cap; the group +// dimension books against the rule's lowest intersecting group. This runs +// alongside the policy-window record so account ceilings accumulate in their own +// windows (commonly monthly) independently of the per-policy window. +func (m *managerImpl) RecordAccountBudgetUsage(ctx context.Context, accountID, userID string, groupIDs []string, tokensIn, tokensOut int64, costUSD float64) error { + if accountID == "" { + return status.Errorf(status.InvalidArgument, "account_id is required") + } + if err := validateUsageDeltas(tokensIn, tokensOut, costUSD); err != nil { + return err + } + rules, err := m.store.GetAccountAgentNetworkBudgetRules(ctx, store.LockingStrengthNone, accountID) + if err != nil { + return fmt.Errorf("list account budget rules: %w", err) + } + set := make(map[types.ConsumptionKey]struct{}) + addAccountBudgetKeys(set, PolicySelectionInput{AccountID: accountID, UserID: userID, GroupIDs: groupIDs}, rules, time.Now().UTC()) + if len(set) == 0 { + return nil + } + return m.store.IncrementAgentNetworkConsumptionBatch(ctx, accountID, keysSlice(set), tokensIn, tokensOut, costUSD) +} + +// RecordUsageInput carries everything RecordUsage books for one served request. +type RecordUsageInput struct { + AccountID string + UserID string + AttributionGroupID string // selected policy's attribution group (policy window) + GroupIDs []string + WindowSeconds int64 // selected policy's window; 0 means no policy cap + TokensIn int64 + TokensOut int64 + CostUSD float64 +} + +// RecordUsage books a served request's usage against every counter it touches — +// the selected policy's per-(user, group) window plus every applicable account +// budget rule's own window — deduplicated and written in a single transaction. +// Two counters that collapse to the same (dimension, window) tuple are booked +// once, so a single request can never double-count against one cap. +func (m *managerImpl) RecordUsage(ctx context.Context, in RecordUsageInput) error { + if in.AccountID == "" { + return status.Errorf(status.InvalidArgument, "account_id is required") + } + if err := validateUsageDeltas(in.TokensIn, in.TokensOut, in.CostUSD); err != nil { + return err + } + now := time.Now().UTC() + set := make(map[types.ConsumptionKey]struct{}) + + // Policy-window dimensions are booked only when a policy cap bound this + // request (window > 0). A zero window means catch-all-allow / no policy cap; + // the account fan-out below still books against the budget rules' windows. + if in.WindowSeconds > 0 { + addLimitKeys(set, in.UserID, in.AttributionGroupID, in.WindowSeconds, now) + } + + rules, err := m.store.GetAccountAgentNetworkBudgetRules(ctx, store.LockingStrengthNone, in.AccountID) + if err != nil { + return fmt.Errorf("list account budget rules: %w", err) + } + addAccountBudgetKeys(set, PolicySelectionInput{AccountID: in.AccountID, UserID: in.UserID, GroupIDs: in.GroupIDs}, rules, now) + + if len(set) == 0 { + return nil + } + return m.store.IncrementAgentNetworkConsumptionBatch(ctx, in.AccountID, keysSlice(set), in.TokensIn, in.TokensOut, in.CostUSD) +} + +// addAccountBudgetKeys adds the (dimension, window) keys a served request books +// against every applicable account budget rule into the dedup set. +func addAccountBudgetKeys(set map[types.ConsumptionKey]struct{}, in PolicySelectionInput, rules []*types.AccountBudgetRule, now time.Time) { + for _, r := range rules { + if r == nil || !r.Enabled || !budgetRuleApplies(r, in) { + continue + } + attrGroup := lowestIntersect(r.TargetGroups, in.GroupIDs) + for _, window := range ruleWindows(r) { + addLimitKeys(set, in.UserID, attrGroup, window, now) + } + } +} + +// keysSlice flattens a ConsumptionKey set into a slice. +func keysSlice(set map[types.ConsumptionKey]struct{}) []types.ConsumptionKey { + keys := make([]types.ConsumptionKey, 0, len(set)) + for k := range set { + keys = append(keys, k) + } + return keys +} + +// ruleWindows returns the distinct enabled window lengths a budget rule books +// against (token window and/or budget window, deduplicated). +func ruleWindows(r *types.AccountBudgetRule) []int64 { + var windows []int64 + if r.Limits.TokenLimit.Enabled && r.Limits.TokenLimit.WindowSeconds > 0 { + windows = append(windows, r.Limits.TokenLimit.WindowSeconds) + } + if r.Limits.BudgetLimit.Enabled && r.Limits.BudgetLimit.WindowSeconds > 0 { + bw := r.Limits.BudgetLimit.WindowSeconds + if len(windows) == 0 || windows[0] != bw { + windows = append(windows, bw) + } + } + return windows +} + +// effectiveWindowSeconds returns the window length the proxy should +// hand back to RecordLLMUsage. When both halves are enabled with +// different windows, token_limit wins (the more common config); when +// only one is enabled that one wins; when neither is enabled the +// returned value is 0 — RecordLLMUsage treats 0 as "no limit +// tracking" and skips the increment, which is the right pass-through +// for catch-all-allow policies with no caps configured. +func effectiveWindowSeconds(p *types.Policy) int64 { + if p.Limits.TokenLimit.Enabled && p.Limits.TokenLimit.WindowSeconds > 0 { + return p.Limits.TokenLimit.WindowSeconds + } + if p.Limits.BudgetLimit.Enabled && p.Limits.BudgetLimit.WindowSeconds > 0 { + return p.Limits.BudgetLimit.WindowSeconds + } + return 0 +} + +// lowestIntersect returns the lowest-by-string-sort element of +// callerGroups ∩ sourceGroups. Empty when the intersection is empty. +// Lowest is deterministic so multi-node selection converges. +func lowestIntersect(sourceGroups, callerGroups []string) string { + if len(sourceGroups) == 0 || len(callerGroups) == 0 { + return "" + } + srcSet := make(map[string]struct{}, len(sourceGroups)) + for _, g := range sourceGroups { + srcSet[g] = struct{}{} + } + var best string + for _, g := range callerGroups { + if _, ok := srcSet[g]; !ok { + continue + } + if best == "" || g < best { + best = g + } + } + return best +} + +func anyGroupMatches(sourceGroups []string, callerSet map[string]struct{}) bool { + for _, g := range sourceGroups { + if _, ok := callerSet[g]; ok { + return true + } + } + return false +} + +// isUncapped reports whether a policy has any enabled cap with a +// positive limit value. Mirrors the eval functions' guards: a policy +// with token_limit.enabled=true but every cap value at 0 still +// counts as uncapped because the eval would query nothing and bind +// nothing. +func isUncapped(p *types.Policy) bool { + tl := p.Limits.TokenLimit + if tl.Enabled && tl.WindowSeconds > 0 && (tl.GroupCap > 0 || tl.UserCap > 0) { + return false + } + bl := p.Limits.BudgetLimit + if bl.Enabled && bl.WindowSeconds > 0 && (bl.GroupCapUsd > 0 || bl.UserCapUsd > 0) { + return false + } + return true +} + +// groupCapTokens returns the policy's group-token cap when the token +// limit is enabled, zero otherwise. Drives the primary "bigger pool +// first" sort. +func groupCapTokens(p *types.Policy) int64 { + if p.Limits.TokenLimit.Enabled { + return p.Limits.TokenLimit.GroupCap + } + return 0 +} + +// groupCapBudgetUsd returns the policy's group-budget cap in USD +// when the budget limit is enabled, zero otherwise. Secondary sort +// key after token group cap so budget-only policies still order +// predictably. +func groupCapBudgetUsd(p *types.Policy) float64 { + if p.Limits.BudgetLimit.Enabled { + return p.Limits.BudgetLimit.GroupCapUsd + } + return 0 +} + +// userCapTokens returns the policy's per-user token cap when the +// token limit is enabled, zero otherwise. Tertiary sort key, used +// when group caps tie or are absent. +func userCapTokens(p *types.Policy) int64 { + if p.Limits.TokenLimit.Enabled { + return p.Limits.TokenLimit.UserCap + } + return 0 +} + +// userCapBudgetUsd returns the policy's per-user budget cap in USD +// when the budget limit is enabled, zero otherwise. Quaternary sort +// key for budget-only policies whose group caps tie or are absent. +func userCapBudgetUsd(p *types.Policy) float64 { + if p.Limits.BudgetLimit.Enabled { + return p.Limits.BudgetLimit.UserCapUsd + } + return 0 +} + +func sliceContains(haystack []string, needle string) bool { + for _, v := range haystack { + if v == needle { + return true + } + } + return false +} + +// mockManager fallback so tests that don't care about selection still +// compile. +func (*mockManager) SelectPolicyForRequest(_ context.Context, _ PolicySelectionInput) (*PolicySelectionResult, error) { + return &PolicySelectionResult{Allow: true}, nil +} diff --git a/management/server/agentnetwork/policyselect_account_realstore_test.go b/management/server/agentnetwork/policyselect_account_realstore_test.go new file mode 100644 index 000000000..a3bfc4769 --- /dev/null +++ b/management/server/agentnetwork/policyselect_account_realstore_test.go @@ -0,0 +1,181 @@ +package agentnetwork + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/store" +) + +// GC-2 no-mock enforcement tests for the account-budget ceiling. They drive the +// real store + real consumption accounting through SelectPolicyForRequest and +// RecordAccountBudgetUsage, asserting min-wins (account binds independently of +// policy), targeting (groups + direct users), and the record fan-out. + +func accountWideUserTokenRule(id string, userCap, window int64) *types.AccountBudgetRule { + r := types.NewAccountBudgetRule(realSelectAccount) + r.ID = id + r.Limits.TokenLimit = types.PolicyTokenLimit{Enabled: true, UserCap: userCap, WindowSeconds: window} + return r +} + +// TestSelectPolicy_RealStore_AccountCeilingBindsEvenWithUncappedPolicy proves +// min-wins: the account user ceiling denies once exhausted even though a +// catch-all-allow (uncapped) policy would otherwise pass the request. The +// account gate runs independently of and ahead of policy selection. +func TestSelectPolicy_RealStore_AccountCeilingBindsEvenWithUncappedPolicy(t *testing.T) { + mgr, s := newRealSelectorMgr(t) + ctx := context.Background() + + // An uncapped (catch-all-allow) policy: enabled token limit, zero caps. + uncapped := capPolicy("pol-open", realSelectAccount, []string{"grp-eng"}, "prov-1", 0, 86_400) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, uncapped)) + + // Account-wide user ceiling of 100 tokens in an hourly window. + require.NoError(t, s.SaveAgentNetworkBudgetRule(ctx, accountWideUserTokenRule("ainbud-1", 100, 3_600))) + + in := PolicySelectionInput{AccountID: realSelectAccount, UserID: "user-1", GroupIDs: []string{"grp-eng"}, ProviderID: "prov-1"} + + // Fresh: account ceiling has headroom, uncapped policy wins. + res, err := mgr.SelectPolicyForRequest(ctx, in) + require.NoError(t, err) + assert.True(t, res.Allow, "fresh account ceiling must allow") + + // Drain the account user ceiling via the fan-out path. + require.NoError(t, mgr.RecordAccountBudgetUsage(ctx, realSelectAccount, "user-1", []string{"grp-eng"}, 100, 0, 0)) + + res, err = mgr.SelectPolicyForRequest(ctx, in) + require.NoError(t, err) + assert.False(t, res.Allow, "account ceiling must deny even though the policy is uncapped (min-wins)") + assert.Equal(t, denyCodeAccountTokenCapExceeded, res.DenyCode, "deny must carry the llm_account.* code") +} + +// TestSelectPolicy_RealStore_AccountGroupCeiling proves a group-targeted rule +// binds the caller's group dimension. +func TestSelectPolicy_RealStore_AccountGroupCeiling(t *testing.T) { + mgr, s := newRealSelectorMgr(t) + ctx := context.Background() + + rule := types.NewAccountBudgetRule(realSelectAccount) + rule.ID = "ainbud-grp" + rule.TargetGroups = []string{"grp-eng"} + rule.Limits.BudgetLimit = types.PolicyBudgetLimit{Enabled: true, GroupCapUsd: 5.0, WindowSeconds: 2_592_000} + require.NoError(t, s.SaveAgentNetworkBudgetRule(ctx, rule)) + + in := PolicySelectionInput{AccountID: realSelectAccount, UserID: "user-1", GroupIDs: []string{"grp-eng"}, ProviderID: "prov-1"} + + res, err := mgr.SelectPolicyForRequest(ctx, in) + require.NoError(t, err) + assert.True(t, res.Allow, "fresh group ceiling must allow") + + require.NoError(t, mgr.RecordAccountBudgetUsage(ctx, realSelectAccount, "user-1", []string{"grp-eng"}, 0, 0, 5.0)) + + res, err = mgr.SelectPolicyForRequest(ctx, in) + require.NoError(t, err) + assert.False(t, res.Allow, "group budget ceiling must deny once spent") + assert.Equal(t, denyCodeAccountBudgetCapExceeded, res.DenyCode, "account budget deny code") +} + +// TestSelectPolicy_RealStore_AccountTargetUsersBindsOnlyThatUser proves a +// TargetUsers rule tightens only the named user, leaving others unbound. +func TestSelectPolicy_RealStore_AccountTargetUsersBindsOnlyThatUser(t *testing.T) { + mgr, s := newRealSelectorMgr(t) + ctx := context.Background() + + rule := types.NewAccountBudgetRule(realSelectAccount) + rule.ID = "ainbud-alice" + rule.TargetUsers = []string{"alice"} + rule.Limits.TokenLimit = types.PolicyTokenLimit{Enabled: true, UserCap: 100, WindowSeconds: 3_600} + require.NoError(t, s.SaveAgentNetworkBudgetRule(ctx, rule)) + + // Record alice's usage to the rule window. + require.NoError(t, mgr.RecordAccountBudgetUsage(ctx, realSelectAccount, "alice", nil, 100, 0, 0)) + + aliceIn := PolicySelectionInput{AccountID: realSelectAccount, UserID: "alice", ProviderID: "prov-1"} + res, err := mgr.SelectPolicyForRequest(ctx, aliceIn) + require.NoError(t, err) + assert.False(t, res.Allow, "alice is bound by the TargetUsers rule and is exhausted") + + bobIn := PolicySelectionInput{AccountID: realSelectAccount, UserID: "bob", ProviderID: "prov-1"} + res, err = mgr.SelectPolicyForRequest(ctx, bobIn) + require.NoError(t, err) + assert.True(t, res.Allow, "bob is not in TargetUsers, so the rule must not bind him") +} + +// TestSelectPolicy_RealStore_AccountRuleRecordsToOwnWindow proves the record +// fan-out books usage in the rule's own window (distinct from any policy +// window), so the account ceiling accumulates independently. +func TestSelectPolicy_RealStore_AccountRuleRecordsToOwnWindow(t *testing.T) { + mgr, s := newRealSelectorMgr(t) + ctx := context.Background() + + require.NoError(t, s.SaveAgentNetworkBudgetRule(ctx, accountWideUserTokenRule("ainbud-w", 100, 3_600))) + + require.NoError(t, mgr.RecordAccountBudgetUsage(ctx, realSelectAccount, "user-1", nil, 60, 0, 0)) + + // Same user, a policy-style daily window must NOT see the account-window + // usage — windows are independent counters. + dailyRow, err := s.GetAgentNetworkConsumption(ctx, store.LockingStrengthNone, realSelectAccount, types.DimensionUser, "user-1", 86_400, types.WindowStart(time.Now().UTC(), 86_400)) + require.NoError(t, err) + assert.Equal(t, int64(0), dailyRow.TokensInput+dailyRow.TokensOutput, "daily window must be untouched by the hourly account-rule record") + + // A second record pushes the hourly account window to its cap → deny. + require.NoError(t, mgr.RecordAccountBudgetUsage(ctx, realSelectAccount, "user-1", nil, 40, 0, 0)) + res, err := mgr.SelectPolicyForRequest(ctx, PolicySelectionInput{AccountID: realSelectAccount, UserID: "user-1", ProviderID: "prov-1"}) + require.NoError(t, err) + assert.False(t, res.Allow, "100 tokens recorded in the rule's hourly window must exhaust the 100-token ceiling") + assert.Equal(t, denyCodeAccountTokenCapExceeded, res.DenyCode, "account token deny code") +} + +// TestRecordUsage_RealStore_BooksPolicyAndAccountWindows proves the batched +// post-flight write books the selected policy's window AND every applicable +// account rule's (independent) window in a single call — the #6 batched-write +// path the proxy's RecordLLMUsage RPC now uses. +func TestRecordUsage_RealStore_BooksPolicyAndAccountWindows(t *testing.T) { + mgr, s := newRealSelectorMgr(t) + ctx := context.Background() + + // Policy: 100-token group cap on a daily window. Account rule: 100-token + // user ceiling on an hourly window — an independent counter. + policy := capPolicy("pol-1", realSelectAccount, []string{"grp-eng"}, "prov-1", 100, 86_400) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, policy)) + require.NoError(t, s.SaveAgentNetworkBudgetRule(ctx, accountWideUserTokenRule("ainbud-1", 100, 3_600))) + + in := PolicySelectionInput{AccountID: realSelectAccount, UserID: "user-1", GroupIDs: []string{"grp-eng"}, ProviderID: "prov-1"} + res, err := mgr.SelectPolicyForRequest(ctx, in) + require.NoError(t, err) + require.True(t, res.Allow) + require.Equal(t, "pol-1", res.SelectedPolicyID) + + // One batched record books the policy window (group + user @86400) and the + // account rule window (user @3600) atomically. + require.NoError(t, mgr.RecordUsage(ctx, RecordUsageInput{ + AccountID: realSelectAccount, + UserID: "user-1", + AttributionGroupID: res.AttributionGroupID, + GroupIDs: []string{"grp-eng"}, + WindowSeconds: res.WindowSeconds, + TokensIn: 100, + })) + + // The next selection denies — the account hourly ceiling binds first. + res, err = mgr.SelectPolicyForRequest(ctx, in) + require.NoError(t, err) + assert.False(t, res.Allow, "usage booked by RecordUsage must enforce on the next request") + + // Prove BOTH windows were booked in the one call via a direct batch read. + now := time.Now().UTC() + userKey := types.ConsumptionKey{Kind: types.DimensionUser, DimID: "user-1", WindowSeconds: 3_600, WindowStartUTC: types.WindowStart(now, 3_600)} + groupKey := types.ConsumptionKey{Kind: types.DimensionGroup, DimID: "grp-eng", WindowSeconds: 86_400, WindowStartUTC: types.WindowStart(now, 86_400)} + rows, err := s.GetAgentNetworkConsumptionBatch(ctx, store.LockingStrengthNone, realSelectAccount, []types.ConsumptionKey{userKey, groupKey}) + require.NoError(t, err) + require.Contains(t, rows, userKey, "account rule user/hourly window booked") + require.Contains(t, rows, groupKey, "policy group/daily window booked") + assert.Equal(t, int64(100), rows[userKey].TokensInput, "account hourly user counter") + assert.Equal(t, int64(100), rows[groupKey].TokensInput, "policy daily group counter") +} diff --git a/management/server/agentnetwork/policyselect_realstore_test.go b/management/server/agentnetwork/policyselect_realstore_test.go new file mode 100644 index 000000000..aac4bd276 --- /dev/null +++ b/management/server/agentnetwork/policyselect_realstore_test.go @@ -0,0 +1,214 @@ +package agentnetwork + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/store" +) + +// This file is the no-mock regression guard for policy limit enforcement. +// policyselect_test.go pins the same behavior through a gomock store with +// explicit call-sequence expectations — brittle precisely where the upcoming +// account-budget work (GC-2) refactors the cap-eval primitive and adds an +// account-level gate. These tests drive the REAL sqlite store + REAL +// consumption accounting and assert observable behavior (allow / deny / +// selection / attribution), not which store methods get called. They must keep +// passing unchanged after GC-2 lands, which is what proves "current behavior is +// not changed." + +const realSelectAccount = "acc-realselect-1" + +// newRealSelectorMgr builds a managerImpl backed by a real sqlite test store. +func newRealSelectorMgr(t *testing.T) (*managerImpl, store.Store) { + t.Helper() + ctx := context.Background() + s, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err, "real sqlite test store must come up") + t.Cleanup(cleanup) + return &managerImpl{store: s}, s +} + +// TestSelectPolicy_RealStore_NoApplicablePolicies pins the pass-through: +// nothing targets the (provider, groups) combination, so the selector allows +// without attribution or consumption tracking. +func TestSelectPolicy_RealStore_NoApplicablePolicies(t *testing.T) { + mgr, _ := newRealSelectorMgr(t) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: realSelectAccount, + UserID: "user-1", + GroupIDs: []string{"grp-x"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.True(t, res.Allow, "no applicable policy must pass through as allow") + assert.Empty(t, res.SelectedPolicyID, "no selection when nothing applies") +} + +// TestSelectPolicy_RealStore_AllowAndLowestGroupAttribution pins the v1 +// attribution rule (lowest intersecting group by string sort) through the +// real store, with a fresh (zero) consumption row. +func TestSelectPolicy_RealStore_AllowAndLowestGroupAttribution(t *testing.T) { + mgr, s := newRealSelectorMgr(t) + ctx := context.Background() + + p := capPolicy("pol-A", realSelectAccount, []string{"grp-zz", "grp-aa", "grp-mm"}, "prov-1", 10_000, 86_400) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, p)) + + res, err := mgr.SelectPolicyForRequest(ctx, PolicySelectionInput{ + AccountID: realSelectAccount, + UserID: "user-1", + GroupIDs: []string{"grp-zz", "grp-aa", "grp-mm"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.True(t, res.Allow, "fresh state under cap must allow") + assert.Equal(t, "pol-A", res.SelectedPolicyID, "only applicable policy must be selected") + assert.Equal(t, "grp-aa", res.AttributionGroupID, "lowest-by-sort intersecting group must win") + assert.Equal(t, int64(86_400), res.WindowSeconds, "selected policy's window must be returned") +} + +// TestSelectPolicy_RealStore_LargerPoolWins_FallsThroughWhenExhausted pins the +// core selection behavior end to end. The two policies bind DISTINCT groups so +// they read separate counters — the only shape where fall-through actually +// yields headroom (policies on the same group share one counter, as +// policyselect_test.go notes). Larger pool wins fresh; after real consumption +// drains the larger group, selection falls through to the smaller; once both +// counters are exhausted the request is denied. +func TestSelectPolicy_RealStore_LargerPoolWins_FallsThroughWhenExhausted(t *testing.T) { + mgr, s := newRealSelectorMgr(t) + ctx := context.Background() + + tight := capPolicy("pol-tight", realSelectAccount, []string{"grp-tight"}, "prov-1", 100, 86_400) + tight.CreatedAt = time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + wide := capPolicy("pol-wide", realSelectAccount, []string{"grp-wide"}, "prov-1", 10_000, 86_400) + wide.CreatedAt = time.Date(2026, 2, 1, 0, 0, 0, 0, time.UTC) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, tight)) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, wide)) + + // Caller is in both groups, so both policies apply with independent counters. + in := PolicySelectionInput{ + AccountID: realSelectAccount, + UserID: "user-1", + GroupIDs: []string{"grp-tight", "grp-wide"}, + ProviderID: "prov-1", + } + + // Fresh: larger pool wins. + res, err := mgr.SelectPolicyForRequest(ctx, in) + require.NoError(t, err) + assert.Equal(t, "pol-wide", res.SelectedPolicyID, "larger pool drains first") + + // Drain only the wide group's counter to its cap. + require.NoError(t, mgr.RecordConsumption(ctx, realSelectAccount, types.DimensionGroup, "grp-wide", 86_400, 10_000, 0, 0)) + + // Wide exhausted, tight's separate counter is fresh → fall through to tight. + res, err = mgr.SelectPolicyForRequest(ctx, in) + require.NoError(t, err) + assert.True(t, res.Allow, "tight pool has its own untouched counter") + assert.Equal(t, "pol-tight", res.SelectedPolicyID, "selection falls through to the smaller pool once the larger is exhausted") + + // Drain the tight group's counter too → both exhausted → deny. + require.NoError(t, mgr.RecordConsumption(ctx, realSelectAccount, types.DimensionGroup, "grp-tight", 86_400, 100, 0, 0)) + res, err = mgr.SelectPolicyForRequest(ctx, in) + require.NoError(t, err) + assert.False(t, res.Allow, "both group counters exhausted must deny") + assert.Equal(t, denyCodeTokenCapExceeded, res.DenyCode, "deny code names the offending cap kind") +} + +// TestSelectPolicy_RealStore_BudgetCapDenies pins budget (USD) enforcement +// through the real store: once recorded cost reaches the cap, deny. +func TestSelectPolicy_RealStore_BudgetCapDenies(t *testing.T) { + mgr, s := newRealSelectorMgr(t) + ctx := context.Background() + + p := &types.Policy{ + ID: "pol-budget", + AccountID: realSelectAccount, + Enabled: true, + SourceGroups: []string{"grp-eng"}, + DestinationProviderIDs: []string{"prov-1"}, + Limits: types.PolicyLimits{ + BudgetLimit: types.PolicyBudgetLimit{ + Enabled: true, + GroupCapUsd: 5.0, + WindowSeconds: 86_400, + }, + }, + CreatedAt: time.Now().UTC(), + } + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, p)) + + in := PolicySelectionInput{ + AccountID: realSelectAccount, + UserID: "user-1", + GroupIDs: []string{"grp-eng"}, + ProviderID: "prov-1", + } + + res, err := mgr.SelectPolicyForRequest(ctx, in) + require.NoError(t, err) + assert.True(t, res.Allow, "fresh budget must allow") + + require.NoError(t, mgr.RecordConsumption(ctx, realSelectAccount, types.DimensionGroup, "grp-eng", 86_400, 0, 0, 5.0)) + + res, err = mgr.SelectPolicyForRequest(ctx, in) + require.NoError(t, err) + assert.False(t, res.Allow, "cost at the cap must deny") + assert.Equal(t, denyCodeBudgetCapExceeded, res.DenyCode, "budget deny code must be surfaced") +} + +// TestSelectPolicy_RealStore_GroupCounterSharedAcrossPolicies pins that two +// policies on the same group+window read one shared consumption counter: usage +// recorded once is visible to both, so exhausting the group budget denies +// regardless of which policy would attribute. +func TestSelectPolicy_RealStore_GroupCounterSharedAcrossPolicies(t *testing.T) { + mgr, s := newRealSelectorMgr(t) + ctx := context.Background() + + a := capPolicy("pol-a", realSelectAccount, []string{"grp-eng"}, "prov-1", 1_000, 86_400) + b := capPolicy("pol-b", realSelectAccount, []string{"grp-eng"}, "prov-1", 1_000, 86_400) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, a)) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, b)) + + in := PolicySelectionInput{ + AccountID: realSelectAccount, + UserID: "user-1", + GroupIDs: []string{"grp-eng"}, + ProviderID: "prov-1", + } + + require.NoError(t, mgr.RecordConsumption(ctx, realSelectAccount, types.DimensionGroup, "grp-eng", 86_400, 1_000, 0, 0)) + + res, err := mgr.SelectPolicyForRequest(ctx, in) + require.NoError(t, err) + assert.False(t, res.Allow, "shared group counter at cap denies both equal policies") + assert.Equal(t, denyCodeTokenCapExceeded, res.DenyCode, "token deny code on the shared counter") +} + +// TestSelectPolicy_RealStore_DisabledPolicyIgnored pins that a disabled policy +// is invisible to selection even when it otherwise matches. +func TestSelectPolicy_RealStore_DisabledPolicyIgnored(t *testing.T) { + mgr, s := newRealSelectorMgr(t) + ctx := context.Background() + + p := capPolicy("pol-disabled", realSelectAccount, []string{"grp-eng"}, "prov-1", 10_000, 86_400) + p.Enabled = false + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, p)) + + res, err := mgr.SelectPolicyForRequest(ctx, PolicySelectionInput{ + AccountID: realSelectAccount, + UserID: "user-1", + GroupIDs: []string{"grp-eng"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.True(t, res.Allow, "no enabled policy applies → pass-through allow") + assert.Empty(t, res.SelectedPolicyID, "disabled policy must not be selected") +} diff --git a/management/server/agentnetwork/policyselect_test.go b/management/server/agentnetwork/policyselect_test.go new file mode 100644 index 000000000..8ac8032db --- /dev/null +++ b/management/server/agentnetwork/policyselect_test.go @@ -0,0 +1,641 @@ +package agentnetwork + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/golang/mock/gomock" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/store" + nbstatus "github.com/netbirdio/netbird/shared/management/status" +) + +func newSelectorMgr(t *testing.T, ctrl *gomock.Controller) (*managerImpl, *store.MockStore) { + t.Helper() + mockStore := store.NewMockStore(ctrl) + // SelectPolicyForRequest evaluates the account-budget ceiling before policy + // selection. These policy-selection tests don't exercise account rules, so + // default to "no rules" — the no-mock policyselect_realstore_test.go covers + // the account gate's behavior end to end. + mockStore.EXPECT(). + GetAccountAgentNetworkBudgetRules(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, nil). + AnyTimes() + return &managerImpl{store: mockStore}, mockStore +} + +type usedKey struct { + kind types.ConsumptionDimension + dimID string + window int64 +} + +// expectConsumptionBatch stubs the batched consumption read to return the +// supplied per-(kind, dim, window) counters, filling each row's window start +// from the actual request keys so it always matches what the selector computed. +// Keys absent from used resolve to zero counters. +func expectConsumptionBatch(mockStore *store.MockStore, used map[usedKey]*types.Consumption) { + mockStore.EXPECT(). + GetAgentNetworkConsumptionBatch(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, _ store.LockingStrength, _ string, keys []types.ConsumptionKey) (map[types.ConsumptionKey]*types.Consumption, error) { + out := make(map[types.ConsumptionKey]*types.Consumption) + for _, k := range keys { + if row, ok := used[usedKey{k.Kind, k.DimID, k.WindowSeconds}]; ok { + rc := *row + rc.WindowStartUTC = k.WindowStartUTC + out[k] = &rc + } + } + return out, nil + }). + AnyTimes() +} + +func capPolicy(id, account string, sourceGroups []string, providerID string, tokenCap int64, windowSec int64) *types.Policy { + return &types.Policy{ + ID: id, + AccountID: account, + Enabled: true, + SourceGroups: sourceGroups, + DestinationProviderIDs: []string{providerID}, + Limits: types.PolicyLimits{ + TokenLimit: types.PolicyTokenLimit{ + Enabled: true, + GroupCap: tokenCap, + WindowSeconds: windowSec, + }, + }, + CreatedAt: time.Now().UTC(), + } +} + +// TestSelectPolicy_NoApplicablePolicies covers the pass-through path: +// llm_router authorisation is upstream of selection; when the +// selector finds no policy targeting the (provider, caller-groups) +// combination, it returns Allow with no attribution and lets the +// request continue without consumption tracking. +func TestSelectPolicy_NoApplicablePolicies(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{}, nil) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + UserID: "user-1", + GroupIDs: []string{"grp-x"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.True(t, res.Allow, "no applicable policies = pass-through allow") + assert.Empty(t, res.SelectedPolicyID, "no selection when nothing applies") +} + +// TestSelectPolicy_AllowWithLowestGroupAttribution proves the v1 +// attribution rule: when the caller's groups intersect a policy's +// source_groups in multiple positions, the selector picks the lowest +// group id by string sort so multi-node selection converges. +func TestSelectPolicy_AllowWithLowestGroupAttribution(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + policy := capPolicy("pol-A", "acc-1", []string{"grp-zz", "grp-aa", "grp-mm"}, "prov-1", 10_000, 86_400) + + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{policy}, nil) + // Fresh: zero consumption across the board. + expectConsumptionBatch(mockStore, nil) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + UserID: "user-1", + GroupIDs: []string{"grp-zz", "grp-aa", "grp-mm"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.True(t, res.Allow) + assert.Equal(t, "pol-A", res.SelectedPolicyID) + assert.Equal(t, "grp-aa", res.AttributionGroupID, + "lowest-by-sort intersection wins so multi-node selection converges") + assert.Equal(t, int64(86_400), res.WindowSeconds) +} + +// TestSelectPolicy_LargerPoolWinsAcrossUsageLevels proves the core +// selection rule: among multiple applicable policies with caps, the +// selector picks the one with the larger absolute pool — at every +// usage level, not just at fresh state. The smaller-pool policy is +// only reached when the larger one is exhausted. This is the +// "drain biggest first" semantic operators expect for layered +// tiers; a fraction-based score would flap between the two as +// soon as one is partially used. +func TestSelectPolicy_LargerPoolWinsAcrossUsageLevels(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + tight := capPolicy("pol-tight", "acc-1", []string{"grp-engineers"}, "prov-1", 100, 86_400) + tight.CreatedAt = time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + wide := capPolicy("pol-wide", "acc-1", []string{"grp-engineers"}, "prov-1", 10_000, 86_400) + wide.CreatedAt = time.Date(2026, 2, 1, 0, 0, 0, 0, time.UTC) + + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{tight, wide}, nil) + + // Both partially used. tight at 50/100 (50% used); wide at + // 50/10000 (0.5% used). Old fraction-based algo would pick wide + // here too — but for the wrong reason ("more relative slack"). + // New algo picks wide because its initial group cap is bigger + // (10000 > 100), and that decision is stable as wide drains. + expectConsumptionBatch(mockStore, map[usedKey]*types.Consumption{ + {types.DimensionGroup, "grp-engineers", 86_400}: {TokensInput: 50}, + }) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + UserID: "user-1", + GroupIDs: []string{"grp-engineers"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.Equal(t, "pol-wide", res.SelectedPolicyID, + "the policy with the bigger initial pool wins — operators expect 'drain the privileged tier first', not load-balance across tiers") +} + +// TestSelectPolicy_StaysOnLargerPoolAfterPartialDrain locks the +// stickiness contract reported by operators: with two policies +// where A has a 200-token group cap and B has 150, the very first +// request goes to A AND every subsequent request continues to land +// on A until A's group cap is exhausted — at which point B becomes +// the only candidate. A fraction-based score would flap to B as +// soon as A had any consumption (B's 1.0 fraction beats A's 0.75) +// even though A still has more absolute headroom; that produced +// confusing per-policy attribution ledger entries and stranded +// A's remaining capacity behind B's exhaustion. +func TestSelectPolicy_StaysOnLargerPoolAfterPartialDrain(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + policyA := capPolicy("pol-A-200", "acc-1", []string{"grp-engineers"}, "prov-1", 200, 86_400) + policyB := capPolicy("pol-B-150", "acc-1", []string{"grp-engineers"}, "prov-1", 150, 86_400) + + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{policyA, policyB}, nil) + + // A is partially drained (50/200 used = 25% used; 75% headroom + // remaining). B is fresh (0/150). The old fraction-based score + // would pick B here (1.0 > 0.75 fraction); the new pool-size + // score sticks with A (200 > 150 absolute cap). + expectConsumptionBatch(mockStore, map[usedKey]*types.Consumption{ + {types.DimensionGroup, "grp-engineers", 86_400}: {TokensInput: 50}, + }) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + UserID: "user-1", + GroupIDs: []string{"grp-engineers"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.Equal(t, "pol-A-200", res.SelectedPolicyID, + "once attribution lands on the bigger pool it must STAY there until exhausted — operators expect 'drain A then B', not 'flip to B as soon as A is touched'") +} + +// TestSelectPolicy_FallsThroughToSmallerPoolWhenLargerExhausted +// proves the second half of the stickiness contract: once the +// larger-pool policy IS exhausted, the smaller one takes over. +// Without this we'd deny on requests the smaller policy is fully +// equipped to serve. +func TestSelectPolicy_FallsThroughToSmallerPoolWhenLargerExhausted(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + policyA := capPolicy("pol-A-200", "acc-1", []string{"grp-engineers"}, "prov-1", 200, 86_400) + // B uses a different window length so it has an INDEPENDENT counter — the + // realistic shape for fall-through. On the SAME (group, window) tuple the + // counter is shared, so A's cap of 200 being reached would also exhaust B's + // 150; independent counters are what let A exhaust while B retains headroom. + policyB := capPolicy("pol-B-150", "acc-1", []string{"grp-engineers"}, "prov-1", 150, 3_600) + + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{policyA, policyB}, nil) + + expectConsumptionBatch(mockStore, map[usedKey]*types.Consumption{ + {types.DimensionGroup, "grp-engineers", 86_400}: {TokensInput: 200}, // A: 200 >= 200 → exhausted + {types.DimensionGroup, "grp-engineers", 3_600}: {TokensInput: 100}, // B: 100 < 150 → headroom + }) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + UserID: "user-1", + GroupIDs: []string{"grp-engineers"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.Equal(t, "pol-B-150", res.SelectedPolicyID, + "once the bigger pool is exhausted, the smaller one must take over — denying when capacity remains would strand B's allowance") +} + +// TestSelectPolicy_TiebreakByLargerGroupPool covers the user-reported +// bug: an admin in two groups (Users + Admins) where Users is bound +// by a smaller-group-cap policy (50 group, 100 user) and Admins is +// bound by a bigger-group-cap policy (100 group, 20 user) MUST get +// attributed to the Admins policy on the first request. +// +// Without this rule, the fresh-state fraction is 1.0 for both and +// the older policy wins by created_at. The first 24-token request +// then drains the shared user counter past Admins's tight 20-token +// user cap, locking Admins out of selection forever. The 100-token +// Admins group pool ends up stranded while requests pile onto the +// 50-token Users pool — the opposite of what the operator intended +// when they put the bigger pool on the privileged group. +func TestSelectPolicy_TiebreakByLargerGroupPool(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + // Policy A: Users group, smaller group pool, looser per-user cap. + policyA := &types.Policy{ + ID: "pol-Users", + AccountID: "acc-1", + Enabled: true, + SourceGroups: []string{"grp-Users"}, + DestinationProviderIDs: []string{"prov-1"}, + Limits: types.PolicyLimits{ + TokenLimit: types.PolicyTokenLimit{ + Enabled: true, GroupCap: 50, UserCap: 100, WindowSeconds: 86_400, + }, + }, + // Older — would win the legacy created_at tiebreak. + CreatedAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC), + } + // Policy B: Admins group, bigger group pool, tighter per-user cap. + policyB := &types.Policy{ + ID: "pol-Admins", + AccountID: "acc-1", + Enabled: true, + SourceGroups: []string{"grp-Admins"}, + DestinationProviderIDs: []string{"prov-1"}, + Limits: types.PolicyLimits{ + TokenLimit: types.PolicyTokenLimit{ + Enabled: true, GroupCap: 100, UserCap: 20, WindowSeconds: 86_400, + }, + }, + CreatedAt: time.Date(2026, 2, 1, 0, 0, 0, 0, time.UTC), + } + + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{policyA, policyB}, nil) + // Fresh state: every cap evaluation reads zero usage. + expectConsumptionBatch(mockStore, nil) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + UserID: "user-1", + GroupIDs: []string{"grp-Users", "grp-Admins"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.Equal(t, "pol-Admins", res.SelectedPolicyID, + "the bigger group pool wins the fresh-state tiebreak — picking Users first would burn the shared user counter past Admins's tight user cap on the very first request and strand the bigger Admins pool") + assert.Equal(t, "grp-Admins", res.AttributionGroupID) +} + +// TestSelectPolicy_TiebreakByCreatedAt proves the deterministic +// final tiebreak: when two applicable policies have the same +// headroom fraction AND the same group cap (so the larger-pool rule +// can't differentiate either), the older policy wins so attribution +// is stable across replays. +func TestSelectPolicy_TiebreakByCreatedAt(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + older := capPolicy("pol-old", "acc-1", []string{"grp-engineers"}, "prov-1", 1_000, 86_400) + older.CreatedAt = time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + newer := capPolicy("pol-new", "acc-1", []string{"grp-engineers"}, "prov-1", 1_000, 86_400) + newer.CreatedAt = time.Date(2026, 3, 1, 0, 0, 0, 0, time.UTC) + + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{newer, older}, nil) + // Both at zero consumption → identical headroom fraction. + expectConsumptionBatch(mockStore, nil) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-engineers"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.Equal(t, "pol-old", res.SelectedPolicyID, + "older policy wins on equal-headroom tiebreak so attribution is stable across replays") +} + +// TestSelectPolicy_DeniesWhenAllExhausted proves the deny envelope: +// when every applicable policy has at least one cap fully exhausted, +// the selector returns Allow=false with the most-recent exhaustion's +// deny code + human reason. The proxy's middleware surfaces this as +// a 403 with the canonical llm_policy.* code. +func TestSelectPolicy_DeniesWhenAllExhausted(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + a := capPolicy("pol-a", "acc-1", []string{"grp-engineers"}, "prov-1", 100, 86_400) + b := capPolicy("pol-b", "acc-1", []string{"grp-engineers"}, "prov-1", 200, 86_400) + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{a, b}, nil) + + // Shared group counter at 200: A (cap 100) and B (cap 200) both exhausted. + expectConsumptionBatch(mockStore, map[usedKey]*types.Consumption{ + {types.DimensionGroup, "grp-engineers", 86_400}: {TokensInput: 200}, + }) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-engineers"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.False(t, res.Allow, "every applicable policy exhausted = deny") + assert.Equal(t, denyCodeTokenCapExceeded, res.DenyCode) + assert.Contains(t, res.DenyReason, "token cap exhausted", + "deny reason must name the exhausted cap kind for operator debugging") +} + +// TestSelectPolicy_UncappedPolicyAlwaysWinsAgainstCapped proves the +// catch-all-allow contract: a policy with NO enabled caps wins +// against any capped policy regardless of how much headroom the +// capped one has, because operators who configure unlimited access +// expect requests to attribute there until they explicitly add caps. +func TestSelectPolicy_UncappedPolicyAlwaysWinsAgainstCapped(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + uncapped := &types.Policy{ + ID: "pol-uncapped", + AccountID: "acc-1", + Enabled: true, + SourceGroups: []string{"grp-engineers"}, + DestinationProviderIDs: []string{"prov-1"}, + // All Limits.*.Enabled = false (zero-value). + CreatedAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC), + } + wide := capPolicy("pol-wide", "acc-1", []string{"grp-engineers"}, "prov-1", 1_000_000, 86_400) + wide.CreatedAt = time.Date(2025, 12, 1, 0, 0, 0, 0, time.UTC) // older than uncapped + + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{uncapped, wide}, nil) + // Only the wide policy reads consumption; uncapped doesn't query + // because it has no enabled caps. + expectConsumptionBatch(mockStore, nil) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-engineers"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.Equal(t, "pol-uncapped", res.SelectedPolicyID, + "a no-caps policy must always win selection — that's how operators express 'unlimited access through this path'") + assert.Equal(t, int64(0), res.WindowSeconds, "no caps configured = WindowSeconds=0 so RecordLLMUsage skips counter writes") +} + +// TestSelectPolicy_DisabledPolicyIgnored proves disabled policies +// don't count toward selection — even when they'd otherwise be the +// best match. Operators disable a policy to take it offline; the +// selector must respect that and route through whatever's left. +func TestSelectPolicy_DisabledPolicyIgnored(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + disabled := capPolicy("pol-disabled", "acc-1", []string{"grp-engineers"}, "prov-1", 1_000_000, 86_400) + disabled.Enabled = false + enabled := capPolicy("pol-enabled", "acc-1", []string{"grp-engineers"}, "prov-1", 100, 86_400) + + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{disabled, enabled}, nil) + expectConsumptionBatch(mockStore, nil) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-engineers"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.Equal(t, "pol-enabled", res.SelectedPolicyID, + "disabled policies must be ignored at selection time") +} + +// TestSelectPolicy_StoreErrorPropagates locks the no-fail-open +// contract: a transient store error must surface to the caller, not +// be silently treated as "no policies = allow". A false allow on the +// hot path would let a request slip past every cap. +func TestSelectPolicy_StoreErrorPropagates(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return(nil, errors.New("boom")) + + _, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + }) + require.Error(t, err, "store errors must surface — never fail open on the hot path") +} + +// TestSelectPolicy_RejectsEmptyAccount is the input-validation guard: +// empty account_id is a programmer error and must surface as +// InvalidArgument, not as a silent zero-result lookup. +func TestSelectPolicy_RejectsEmptyAccount(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, _ := newSelectorMgr(t, ctrl) + + _, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{}) + require.Error(t, err) + var sErr *nbstatus.Error + require.True(t, errors.As(err, &sErr)) + assert.Equal(t, nbstatus.InvalidArgument, sErr.Type()) +} + +// TestSelectPolicy_SharesGroupCounterAcrossPolicies locks the +// counter-keying design fork: counters are keyed on (account, +// dim_kind, dim_id, window_hours, window_start) — NOT on policy_id. +// Two policies that target the same group with the SAME window length +// share one bucket: spend booked under policy A is visible to policy +// B's headroom calculation and counts toward B's cap. +// +// This is what makes "operator's per-group enforcement" sane — caps +// describe how much a GROUP can use, not how much each policy owes. +func TestSelectPolicy_SharesGroupCounterAcrossPolicies(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + // Two policies, both targeting grp-engineers + prov-1, same 24h + // window length. Different cap sizes. + policyA := capPolicy("pol-A", "acc-1", []string{"grp-engineers"}, "prov-1", 1_000, 86_400) + policyB := capPolicy("pol-B", "acc-1", []string{"grp-engineers"}, "prov-1", 5_000, 86_400) + + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{policyA, policyB}, nil) + // Both policies query the SAME consumption row — same dim_id, + // same window_hours, same window_start. The mock returns the + // same row for both calls, simulating the shared counter. + expectConsumptionBatch(mockStore, map[usedKey]*types.Consumption{ + {types.DimensionGroup, "grp-engineers", 86_400}: {TokensInput: 800}, + }) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-engineers"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + // 800 used → policy A has 200 tokens left of 1000 (20% headroom); + // policy B has 4200 left of 5000 (84% headroom). B wins. + assert.Equal(t, "pol-B", res.SelectedPolicyID, + "the SAME 800 tokens count toward both policies — counters share the (group, window) key, caps differ per policy") +} + +// TestSelectPolicy_AntiFallThroughOnLowestGroup locks the no-fall- +// through behaviour: when a caller is in multiple of a policy's +// source_groups and the lowest-by-sort group is exhausted, we DENY +// rather than fall through to a less-loaded sibling. Per-group caps +// are independent (each group has its own bucket), but attribution +// is one-shot — operators wanting fall-through must split into +// separate policies. +// +// This nails down semantics future contributors might "improve" into +// fall-through behaviour by accident. +func TestSelectPolicy_AntiFallThroughOnLowestGroup(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + // Policy targets two groups; caller is in both. + policy := capPolicy("pol-1", "acc-1", []string{"grp-aaa", "grp-bbb"}, "prov-1", 100, 86_400) + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{policy}, nil) + + // grp-aaa is the lowest by sort → attribution picks it, and the + // prefetch only collects the attribution group's key. We exhaust + // grp-aaa (100/100); grp-bbb's counter is never requested because the + // selector attributes one-shot to the lowest group, so it can't fall + // through to a less-loaded sibling. + expectConsumptionBatch(mockStore, map[usedKey]*types.Consumption{ + {types.DimensionGroup, "grp-aaa", 86_400}: {TokensInput: 100}, + }) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-aaa", "grp-bbb"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.False(t, res.Allow, + "lowest-group-by-sort attribution does NOT fall through to a less-loaded sibling — operators wanting fall-through must split into separate policies") + assert.Equal(t, denyCodeTokenCapExceeded, res.DenyCode) + assert.Contains(t, res.DenyReason, "pol-1", + "deny reason names the exhausted policy id so operators can grep it from the access log") +} + +// TestSelectPolicy_BudgetOnlyExhaustionDenies covers the symmetric +// path to TestSelectPolicy_DeniesWhenAllExhausted but for the budget +// cap: a policy with token_limit DISABLED and budget_limit at-cap +// must deny with llm_policy.budget_cap_exceeded (not the token code). +// +// Without this, the budget evaluation path in evalBudgetCap could +// silently regress and we'd still pass DeniesWhenAllExhausted (which +// only exercises tokens). +func TestSelectPolicy_BudgetOnlyExhaustionDenies(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + policy := &types.Policy{ + ID: "pol-budget", + AccountID: "acc-1", + Enabled: true, + SourceGroups: []string{"grp-engineers"}, + DestinationProviderIDs: []string{"prov-1"}, + Limits: types.PolicyLimits{ + TokenLimit: types.PolicyTokenLimit{Enabled: false}, + BudgetLimit: types.PolicyBudgetLimit{ + Enabled: true, + GroupCapUsd: 10.00, + WindowSeconds: 86_400, + }, + }, + CreatedAt: time.Now().UTC(), + } + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{policy}, nil) + expectConsumptionBatch(mockStore, map[usedKey]*types.Consumption{ + {types.DimensionGroup, "grp-engineers", 86_400}: {CostUSD: 10.50}, // over the $10 cap + }) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-engineers"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.False(t, res.Allow, "budget cap exhausted must deny independently of any token cap state") + assert.Equal(t, denyCodeBudgetCapExceeded, res.DenyCode, + "deny code must be the budget code — token-only deny would silently regress the budget evaluation path") + assert.Contains(t, res.DenyReason, "budget", "deny reason names the budget cap kind for operator debugging") +} + +// TestSelectPolicy_BudgetTighterThanTokenWins is the dual-cap headroom +// fork: when both Token and Budget are enabled on the same policy, +// the SMALLER remaining ratio gates the policy. A policy with +// abundant token headroom but near-zero budget headroom must deny on +// budget, not pass on tokens. +func TestSelectPolicy_BudgetTighterThanTokenWins(t *testing.T) { + ctrl := gomock.NewController(t) + mgr, mockStore := newSelectorMgr(t, ctrl) + + policy := &types.Policy{ + ID: "pol-dual", + AccountID: "acc-1", + Enabled: true, + SourceGroups: []string{"grp-engineers"}, + DestinationProviderIDs: []string{"prov-1"}, + Limits: types.PolicyLimits{ + TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 10_000_000, WindowSeconds: 86_400}, + BudgetLimit: types.PolicyBudgetLimit{Enabled: true, GroupCapUsd: 1.00, WindowSeconds: 86_400}, + }, + CreatedAt: time.Now().UTC(), + } + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), "acc-1"). + Return([]*types.Policy{policy}, nil) + // One shared counter carries both token usage (ample headroom) and cost + // (at the $1 budget cap); the tighter budget cap gates the policy. + expectConsumptionBatch(mockStore, map[usedKey]*types.Consumption{ + {types.DimensionGroup, "grp-engineers", 86_400}: {TokensInput: 100, CostUSD: 1.00}, + }) + + res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{ + AccountID: "acc-1", + GroupIDs: []string{"grp-engineers"}, + ProviderID: "prov-1", + }) + require.NoError(t, err) + assert.False(t, res.Allow, + "the tighter of (token, budget) wins — abundant token headroom must NOT mask an exhausted budget") + assert.Equal(t, denyCodeBudgetCapExceeded, res.DenyCode) +} diff --git a/management/server/agentnetwork/reconcile.go b/management/server/agentnetwork/reconcile.go new file mode 100644 index 000000000..319553ebc --- /dev/null +++ b/management/server/agentnetwork/reconcile.go @@ -0,0 +1,131 @@ +package agentnetwork + +import ( + "context" + + log "github.com/sirupsen/logrus" + + rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" + "github.com/netbirdio/netbird/management/server/types" + "github.com/netbirdio/netbird/shared/management/proto" +) + +// reconcile recomputes the synthesised reverse-proxy services for an +// account, diffs them against the previously-synthesised set in the +// in-memory cache, and emits Create / Update / Delete proxy mappings +// to the affected clusters. Also triggers a peer-side network-map +// recompute via accountManager.UpdateAccountPeers so the +// private-service ACL injection picks up the new state immediately. +// +// Reconcile failures are logged and swallowed — the underlying CRUD +// has already completed, and the next mutation (or proxy reconnect) +// will re-converge the cluster's view. +func (m *managerImpl) reconcile(ctx context.Context, accountID string) { + if accountID == "" { + return + } + + defer func() { + if m.accountManager != nil { + m.accountManager.UpdateAccountPeers(ctx, accountID, types.UpdateReason{ + Resource: types.UpdateResourceService, + Operation: types.UpdateOperationUpdate, + }) + } + }() + + if m.proxyController == nil { + return + } + + services, err := SynthesizeServices(ctx, m.store, accountID) + if err != nil { + log.WithContext(ctx).WithError(err).Warnf("agent-network reconcile: synthesise services for account %s", accountID) + return + } + + oidcCfg := m.proxyController.GetOIDCValidationConfig() + current := make(map[string]*proto.ProxyMapping, len(services)) + for _, svc := range services { + if svc == nil || svc.ID == "" { + continue + } + current[svc.ID] = svc.ToProtoMapping(rpservice.Update, "", oidcCfg) + } + + m.reconcileMu.Lock() + previous := m.reconcileCache[accountID] + if previous == nil { + previous = make(map[string]*proto.ProxyMapping) + } + + creates, updates, deletes := diffMappings(previous, current) + if len(current) == 0 { + delete(m.reconcileCache, accountID) + } else { + m.reconcileCache[accountID] = current + } + m.reconcileMu.Unlock() + + for _, mapping := range creates { + mapping.Type = proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED + m.proxyController.SendServiceUpdateToCluster(ctx, accountID, mapping, clusterFromMapping(mapping)) + } + for _, mapping := range updates { + mapping.Type = proto.ProxyMappingUpdateType_UPDATE_TYPE_MODIFIED + m.proxyController.SendServiceUpdateToCluster(ctx, accountID, mapping, clusterFromMapping(mapping)) + } + for _, mapping := range deletes { + mapping.Type = proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED + m.proxyController.SendServiceUpdateToCluster(ctx, accountID, mapping, clusterFromMapping(mapping)) + } +} + +// diffMappings classifies the previous→current transition for a +// single account into Create / Update / Delete sets. +// +// Cluster moves (current.cluster != previous.cluster) are surfaced as +// a Delete on the old cluster + Create on the new — handled by +// emitting both a delete (on previous mapping) and a create (on the +// current mapping) for that service ID. +func diffMappings(previous, current map[string]*proto.ProxyMapping) (creates, updates, deletes []*proto.ProxyMapping) { + for id, cur := range current { + prev, existed := previous[id] + switch { + case !existed: + creates = append(creates, cur) + case prev.GetDomain() == "" || cur.GetAccountId() == prev.GetAccountId() && currentClusterChanged(prev, cur): + deletes = append(deletes, prev) + creates = append(creates, cur) + default: + updates = append(updates, cur) + } + } + for id, prev := range previous { + if _, stillThere := current[id]; !stillThere { + deletes = append(deletes, prev) + } + } + return creates, updates, deletes +} + +func currentClusterChanged(prev, cur *proto.ProxyMapping) bool { + return clusterFromMapping(prev) != clusterFromMapping(cur) +} + +// clusterFromMapping returns the cluster the mapping should be sent +// to. ProxyMapping doesn't carry the cluster directly, so we rely on +// the synthesised service's domain (`.`) and split on +// the first '.'. +func clusterFromMapping(m *proto.ProxyMapping) string { + if m == nil { + return "" + } + domain := m.GetDomain() + for i := 0; i < len(domain); i++ { + if domain[i] == '.' { + return domain[i+1:] + } + } + return "" +} diff --git a/management/server/agentnetwork/reconcile_test.go b/management/server/agentnetwork/reconcile_test.go new file mode 100644 index 000000000..5a92b60a0 --- /dev/null +++ b/management/server/agentnetwork/reconcile_test.go @@ -0,0 +1,232 @@ +package agentnetwork + +import ( + "context" + "testing" + + "github.com/golang/mock/gomock" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy" + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/store" + "github.com/netbirdio/netbird/shared/management/proto" +) + +func newReconcileMgr(t *testing.T, ctrl *gomock.Controller) (*managerImpl, *store.MockStore, *proxy.MockController) { + t.Helper() + mockStore := store.NewMockStore(ctrl) + mockProxy := proxy.NewMockController(ctrl) + return &managerImpl{ + store: mockStore, + proxyController: mockProxy, + reconcileCache: make(map[string]map[string]*proto.ProxyMapping), + }, mockStore, mockProxy +} + +func newReconcileTestProvider() *types.Provider { + return &types.Provider{ + ID: "prov-1", + AccountID: "acct-1", + ProviderID: "openai_api", + Name: "OpenAI", + UpstreamURL: "https://api.openai.com", + APIKey: "sk-test-key", + Enabled: true, + SessionPrivateKey: "test-priv-key", + SessionPublicKey: "test-pub-key", + } +} + +func newReconcileTestPolicy(providerID, sourceGroupID string) *types.Policy { + return &types.Policy{ + ID: "pol-1", + AccountID: "acct-1", + Name: "engineers", + Enabled: true, + SourceGroups: []string{sourceGroupID}, + DestinationProviderIDs: []string{providerID}, + } +} + +func newReconcileTestSettings() *types.Settings { + return &types.Settings{ + AccountID: "acct-1", + Cluster: "eu.proxy.netbird.io", + Subdomain: "violet", + } +} + +func expectReconcileSynthInputs(mockStore *store.MockStore, ctx context.Context, providers []*types.Provider, policies []*types.Policy, guardrails []*types.Guardrail) { + mockStore.EXPECT(). + GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "acct-1"). + Return(newReconcileTestSettings(), nil) + mockStore.EXPECT(). + GetAccountAgentNetworkProviders(ctx, store.LockingStrengthNone, "acct-1"). + Return(providers, nil) + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(ctx, store.LockingStrengthNone, "acct-1"). + Return(policies, nil) + mockStore.EXPECT(). + GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, "acct-1"). + Return(guardrails, nil) +} + +func TestReconcile_FirstSynth_EmitsCreate(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mgr, mockStore, mockProxy := newReconcileMgr(t, ctrl) + provider := newReconcileTestProvider() + policy := newReconcileTestPolicy(provider.ID, "grp-eng") + + expectReconcileSynthInputs(mockStore, ctx, []*types.Provider{provider}, []*types.Policy{policy}, []*types.Guardrail{}) + mockProxy.EXPECT().GetOIDCValidationConfig().Return(proxy.OIDCValidationConfig{}) + + var sentMappings []*proto.ProxyMapping + mockProxy.EXPECT(). + SendServiceUpdateToCluster(ctx, "acct-1", gomock.Any(), "eu.proxy.netbird.io"). + Do(func(_ context.Context, _ string, m *proto.ProxyMapping, _ string) { + sentMappings = append(sentMappings, m) + }) + + mgr.reconcile(ctx, "acct-1") + + require.Len(t, sentMappings, 1, "first synth must emit one mapping") + assert.Equal(t, proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED, sentMappings[0].Type, "first synth is a Create") + assert.Equal(t, "agent-net-svc-acct-1", sentMappings[0].Id, "stable account-scoped virtual service id") + assert.Equal(t, "violet.eu.proxy.netbird.io", sentMappings[0].Domain, "domain comes from settings (subdomain.cluster)") + + mgr.reconcileMu.Lock() + cached := mgr.reconcileCache["acct-1"] + mgr.reconcileMu.Unlock() + require.Len(t, cached, 1, "cache must hold the synth result for next diff") +} + +func TestReconcile_NoChange_EmitsNothingExtra(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mgr, mockStore, mockProxy := newReconcileMgr(t, ctrl) + provider := newReconcileTestProvider() + policy := newReconcileTestPolicy(provider.ID, "grp-eng") + + // Two identical synth runs. + mockStore.EXPECT(). + GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "acct-1"). + Return(newReconcileTestSettings(), nil).Times(2) + mockStore.EXPECT(). + GetAccountAgentNetworkProviders(ctx, store.LockingStrengthNone, "acct-1"). + Return([]*types.Provider{provider}, nil).Times(2) + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(ctx, store.LockingStrengthNone, "acct-1"). + Return([]*types.Policy{policy}, nil).Times(2) + mockStore.EXPECT(). + GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, "acct-1"). + Return([]*types.Guardrail{}, nil).Times(2) + mockProxy.EXPECT().GetOIDCValidationConfig().Return(proxy.OIDCValidationConfig{}).Times(2) + + createCalls := 0 + updateCalls := 0 + mockProxy.EXPECT(). + SendServiceUpdateToCluster(ctx, "acct-1", gomock.Any(), gomock.Any()). + Do(func(_ context.Context, _ string, m *proto.ProxyMapping, _ string) { + switch m.Type { + case proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED: + createCalls++ + case proto.ProxyMappingUpdateType_UPDATE_TYPE_MODIFIED: + updateCalls++ + } + }). + AnyTimes() + + mgr.reconcile(ctx, "acct-1") + mgr.reconcile(ctx, "acct-1") + + assert.Equal(t, 1, createCalls, "first reconcile creates") + assert.Equal(t, 1, updateCalls, "second reconcile re-pushes as Modified (no semantic change but mapping fields refresh)") +} + +func TestReconcile_PolicyRemoved_EmitsDelete(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mgr, mockStore, mockProxy := newReconcileMgr(t, ctrl) + provider := newReconcileTestProvider() + policy := newReconcileTestPolicy(provider.ID, "grp-eng") + + gomock.InOrder( + // First reconcile: provider + policy, synthesised. + mockStore.EXPECT().GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "acct-1").Return(newReconcileTestSettings(), nil), + mockStore.EXPECT().GetAccountAgentNetworkProviders(ctx, store.LockingStrengthNone, "acct-1").Return([]*types.Provider{provider}, nil), + mockStore.EXPECT().GetAccountAgentNetworkPolicies(ctx, store.LockingStrengthNone, "acct-1").Return([]*types.Policy{policy}, nil), + mockStore.EXPECT().GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, "acct-1").Return([]*types.Guardrail{}, nil), + // Second reconcile: policy gone, provider stays but no longer referenced. + mockStore.EXPECT().GetAgentNetworkSettings(ctx, store.LockingStrengthNone, "acct-1").Return(newReconcileTestSettings(), nil), + mockStore.EXPECT().GetAccountAgentNetworkProviders(ctx, store.LockingStrengthNone, "acct-1").Return([]*types.Provider{provider}, nil), + mockStore.EXPECT().GetAccountAgentNetworkPolicies(ctx, store.LockingStrengthNone, "acct-1").Return([]*types.Policy{}, nil), + ) + mockProxy.EXPECT().GetOIDCValidationConfig().Return(proxy.OIDCValidationConfig{}).AnyTimes() + + var seenTypes []proto.ProxyMappingUpdateType + mockProxy.EXPECT(). + SendServiceUpdateToCluster(ctx, "acct-1", gomock.Any(), "eu.proxy.netbird.io"). + Do(func(_ context.Context, _ string, m *proto.ProxyMapping, _ string) { + seenTypes = append(seenTypes, m.Type) + }). + AnyTimes() + + mgr.reconcile(ctx, "acct-1") + mgr.reconcile(ctx, "acct-1") + + require.Len(t, seenTypes, 2, "create then delete") + assert.Equal(t, proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED, seenTypes[0]) + assert.Equal(t, proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED, seenTypes[1]) + + mgr.reconcileMu.Lock() + _, present := mgr.reconcileCache["acct-1"] + mgr.reconcileMu.Unlock() + assert.False(t, present, "cache for the account must be cleared once nothing is synthesised") +} + +func TestReconcile_NilProxyController_NoOp(t *testing.T) { + ctx := context.Background() + mgr := &managerImpl{ + reconcileCache: make(map[string]map[string]*proto.ProxyMapping), + } + // Must not panic; must not query the store. + mgr.reconcile(ctx, "acct-1") +} + +func TestReconcile_EmptyAccountID_NoOp(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mgr, _, _ := newReconcileMgr(t, ctrl) + // Empty accountID short-circuits before any store call. + mgr.reconcile(ctx, "") +} + +func TestClusterFromMapping(t *testing.T) { + tests := []struct { + name string + domain string + want string + }{ + {"simple", "openai.eu.proxy.netbird.io", "eu.proxy.netbird.io"}, + {"deeply nested", "a.b.c.d", "b.c.d"}, + {"no dot", "openai", ""}, + {"empty", "", ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := clusterFromMapping(&proto.ProxyMapping{Domain: tt.domain}) + assert.Equal(t, tt.want, got) + }) + } +} diff --git a/management/server/agentnetwork/synthesizer.go b/management/server/agentnetwork/synthesizer.go new file mode 100644 index 000000000..156b41df3 --- /dev/null +++ b/management/server/agentnetwork/synthesizer.go @@ -0,0 +1,1059 @@ +package agentnetwork + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/url" + "sort" + "strings" + + rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" + "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey" + "github.com/netbirdio/netbird/management/server/agentnetwork/catalog" + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/store" + "github.com/netbirdio/netbird/shared/management/status" +) + +// apiKeyPlaceholder is the literal substituted with the provider's +// decrypted API key in catalog AuthHeaderTemplate strings. +const apiKeyPlaceholder = "${API_KEY}" //nolint:gosec // template marker, not a credential + +// gcpKeyfilePrefix marks an api_key that holds a base64-encoded GCP +// service-account JSON key ("keyfile::") rather than a static bearer +// token; the proxy mints OAuth tokens from it. Mirrors Aperture's convention. +const gcpKeyfilePrefix = "keyfile::" + +// SynthesizedServiceIDPrefix prefixes the in-memory ID of every +// reverse-proxy service synthesised from Agent Network state. One +// synthesised service exists per (account, cluster); the suffix is the +// account ID so the proxy can dedup mappings cleanly. +const SynthesizedServiceIDPrefix = "agent-net-svc-" + +// agentNetworkRequestCaptureBytes is the request-side body capture cap. +// Kept modest: oversized requests (a long conversation's context can be +// many MB) have their routing fields recovered by the proxy's tolerant +// scan rather than buffered here, so there's no need to size this to the +// largest possible request. +const agentNetworkRequestCaptureBytes = 1 << 20 + +// agentNetworkResponseCaptureBytes is the response-side body capture cap. +// Token usage lives in the trailing SSE message_delta event, so the +// captured prefix must reach the end of the stream. Unlike a request's +// unbounded context, a single response is hard-capped by the model's max +// output tokens (~128K on Opus → a few hundred KB of gzipped SSE even +// with thinking), so 8 MiB is comfortably above any real response and is +// effectively unlimited here — not a moving ceiling. The proxy clamps to +// its own MaxBodyCapBytes at apply time. +const agentNetworkResponseCaptureBytes = 8 << 20 + +// agentNetworkCaptureContentTypes is the set of content types whose +// bodies the proxy buffers for the LLM middlewares. JSON covers +// buffered request and response bodies; SSE covers streaming +// responses (the response parser sums delta tokens across chunks). +var agentNetworkCaptureContentTypes = []string{ + "application/json", + "text/event-stream", +} + +// Middleware IDs the synthesised target chain registers, mirroring the +// proxy-side built-in registry. Order matters: on_request runs in the +// order they're listed; on_response runs in reverse, so cost_meter must +// come BEFORE llm_response_parser in the slice so the parser populates +// tokens before the cost meter reads them. +const ( + middlewareIDLLMRequestParser = "llm_request_parser" + middlewareIDLLMRouter = "llm_router" + middlewareIDLLMIdentityInject = "llm_identity_inject" + middlewareIDLLMLimitCheck = "llm_limit_check" + middlewareIDLLMGuardrail = "llm_guardrail" + middlewareIDCostMeter = "cost_meter" + middlewareIDLLMResponseParser = "llm_response_parser" + middlewareIDLLMLimitRecord = "llm_limit_record" +) + +// SynthesizeServicesForCluster walks every account's agent-network +// settings row pinned to clusterAddr and synthesises the per-account +// gateway service. Used by the proxy-mapping snapshot path where the +// connecting proxy has a specific cluster address and cares about every +// account that routes through it. +// +// Returns nil (no error) when no settings row references the cluster. +// Per-account synthesis failures are skipped rather than dropping every +// account on the cluster. +func SynthesizeServicesForCluster(ctx context.Context, s store.Store, clusterAddr string) ([]*rpservice.Service, error) { + clusterAddr = strings.TrimSpace(clusterAddr) + if clusterAddr == "" { + return nil, nil + } + + settingsRows, err := s.GetAgentNetworkSettingsByCluster(ctx, store.LockingStrengthNone, clusterAddr) + if err != nil { + return nil, fmt.Errorf("list agent network settings on cluster: %w", err) + } + if len(settingsRows) == 0 { + return nil, nil + } + + var out []*rpservice.Service + for _, settings := range settingsRows { + if settings == nil { + continue + } + services, serr := SynthesizeServices(ctx, s, settings.AccountID) + if serr != nil { + continue + } + for _, svc := range services { + if svc != nil && svc.ProxyCluster == clusterAddr { + out = append(out, svc) + } + } + } + return out, nil +} + +// SynthesizeServiceForDomain resolves a single agent-network service by its +// public endpoint domain. It lists the (few) settings rows on the domain's +// cluster, matches the one whose endpoint equals the domain, and synthesises +// only that account — avoiding full per-account synthesis for every tenant on +// the cluster, which is what auth/session paths previously paid. Returns nil +// (no error) when no account owns the domain. +func SynthesizeServiceForDomain(ctx context.Context, s store.Store, domain string) (*rpservice.Service, error) { + domain = strings.TrimSpace(domain) + cluster := clusterFromDomain(domain) + if domain != "" && cluster != "" { + settingsRows, err := s.GetAgentNetworkSettingsByCluster(ctx, store.LockingStrengthNone, cluster) + if err != nil { + return nil, fmt.Errorf("list agent network settings on cluster: %w", err) + } + for _, settings := range settingsRows { + if settings == nil || settings.Endpoint() != domain { + continue + } + services, serr := SynthesizeServices(ctx, s, settings.AccountID) + if serr != nil { + return nil, serr + } + for _, svc := range services { + if svc != nil && svc.Domain == domain { + return svc, nil + } + } + break + } + } + return nil, nil //nolint:nilnil // optional lookup: no account owns the domain +} + +// clusterFromDomain returns the cluster portion of an endpoint domain (every +// label after the first). +func clusterFromDomain(domain string) string { + if i := strings.IndexByte(domain, '.'); i >= 0 { + return domain[i+1:] + } + return "" +} + +// SynthesizeServices builds the in-memory reverse-proxy service that +// fronts the account's agent-network gateway. Returns nil when the +// account has no settings row, no enabled providers, or no enabled +// policies — in any of those cases there's nothing useful to expose. +// +// One service per (account, settings.Cluster) is emitted. The router +// middleware encodes a denormalised model→provider routing table +// (auth headers + decrypted API keys baked in); the policy_check +// middleware encodes per-provider authorised group IDs derived from +// the account's enabled policies. +// +// Services are NEVER persisted — callers regenerate them on every +// network-map / proxy-mapping cycle from current state. +func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([]*rpservice.Service, error) { + settings, ok, err := loadSettings(ctx, s, accountID) + if err != nil { + return nil, err + } + if !ok || strings.TrimSpace(settings.Cluster) == "" { + return nil, nil + } + + providers, err := s.GetAccountAgentNetworkProviders(ctx, store.LockingStrengthNone, accountID) + if err != nil { + return nil, fmt.Errorf("list agent network providers: %w", err) + } + enabledProviders := filterEnabledProviders(providers) + if len(enabledProviders) == 0 { + return nil, nil + } + + policies, err := s.GetAccountAgentNetworkPolicies(ctx, store.LockingStrengthNone, accountID) + if err != nil { + return nil, fmt.Errorf("list agent network policies: %w", err) + } + enabledPolicies := filterEnabledPolicies(policies) + if len(enabledPolicies) == 0 { + return nil, nil + } + + // Backfill any missing session keypairs before deriving a + // service-level keypair. Old rows pre-date the column; treating + // the gap as a no-op produces an immediate dial failure, so we + // fix it once here and persist for future cycles. + for _, p := range enabledProviders { + if p.SessionPrivateKey != "" && p.SessionPublicKey != "" { + continue + } + if err := backfillProviderSessionKeys(ctx, s, p); err != nil { + return nil, fmt.Errorf("backfill session keys for provider %s: %w", p.ID, err) + } + } + + guardrails, err := s.GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, accountID) + if err != nil { + return nil, fmt.Errorf("list agent network guardrails: %w", err) + } + guardrailsByID := make(map[string]*types.Guardrail, len(guardrails)) + for _, g := range guardrails { + if g != nil { + guardrailsByID[g.ID] = g + } + } + + groupIndex := indexProviderGroups(enabledPolicies) + + routerCfgJSON, err := buildRouterConfigJSON(enabledProviders, groupIndex) + if err != nil { + return nil, err + } + + identityInjectJSON, err := buildIdentityInjectConfigJSON(enabledProviders, groupIndex) + if err != nil { + return nil, err + } + + mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID) + applyAccountCollectionControls(&mergedGuardrails, settings) + guardrailJSON, err := marshalGuardrailConfig(mergedGuardrails) + if err != nil { + return nil, err + } + + // Use the merged decision (account settings OR policy-required redaction), + // not the raw account flag, so a policy that mandates PII redaction is + // honored by the capture parsers even when the account toggle is off. + middlewares := buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON, mergedGuardrails.PromptCapture.RedactPii, mergedGuardrails.PromptCapture.Enabled) + + priv, pub, err := pickServiceSessionKeys(enabledProviders) + if err != nil { + return nil, err + } + + svc := buildAccountService(accountID, settings, enabledPolicies, middlewares, priv, pub) + return []*rpservice.Service{svc}, nil +} + +// loadSettings returns the account's agent-network settings row. The +// boolean reports whether a row exists; a status.NotFound surfaces as +// (nil, false, nil) so callers can treat "no settings" as "no +// synthesis" without inspecting error types themselves. +func loadSettings(ctx context.Context, s store.Store, accountID string) (*types.Settings, bool, error) { + settings, err := s.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID) + if err == nil { + return settings, true, nil + } + var sErr *status.Error + if errors.As(err, &sErr) && sErr.Type() == status.NotFound { + return nil, false, nil + } + return nil, false, fmt.Errorf("get agent network settings: %w", err) +} + +// filterEnabledProviders returns the subset of enabled providers, sorted +// by created_at ascending so the router config is deterministic and +// first-match-wins is stable across synthesis cycles. +func filterEnabledProviders(providers []*types.Provider) []*types.Provider { + out := make([]*types.Provider, 0, len(providers)) + for _, p := range providers { + if p == nil || !p.Enabled { + continue + } + out = append(out, p) + } + sort.SliceStable(out, func(i, j int) bool { + if !out[i].CreatedAt.Equal(out[j].CreatedAt) { + return out[i].CreatedAt.Before(out[j].CreatedAt) + } + return out[i].ID < out[j].ID + }) + return out +} + +// filterEnabledPolicies returns the subset of enabled policies. +func filterEnabledPolicies(policies []*types.Policy) []*types.Policy { + out := make([]*types.Policy, 0, len(policies)) + for _, p := range policies { + if p == nil || !p.Enabled { + continue + } + out = append(out, p) + } + return out +} + +// backfillProviderSessionKeys mints an ed25519 session keypair on a +// provider row that doesn't have one yet (rows created before the +// keys were persistent fields) and persists it via the store so +// subsequent cycles get stable keys. +func backfillProviderSessionKeys(ctx context.Context, s store.Store, p *types.Provider) error { + pair, err := sessionkey.GenerateKeyPair() + if err != nil { + return fmt.Errorf("generate session keys for provider %s: %w", p.ID, err) + } + p.SessionPrivateKey = pair.PrivateKey + p.SessionPublicKey = pair.PublicKey + if err := s.SaveAgentNetworkProvider(ctx, p); err != nil { + return fmt.Errorf("persist backfilled session keys for provider %s: %w", p.ID, err) + } + return nil +} + +// pickServiceSessionKeys returns the keypair the synthesised gateway +// service signs / verifies session JWTs with. The PoC reuses the first +// enabled provider's keypair so existing session cookies survive +// provider edits as long as the first-by-created_at provider stays in +// place. Returns an error when no provider has a usable keypair after +// backfill — that surfaces a misconfigured account loudly instead of +// emitting a service the proxy will reject as "invalid session public +// key size". +func pickServiceSessionKeys(providers []*types.Provider) (priv, pub string, err error) { + for _, p := range providers { + if p.SessionPrivateKey != "" && p.SessionPublicKey != "" { + return p.SessionPrivateKey, p.SessionPublicKey, nil + } + } + return "", "", fmt.Errorf("no provider with session keypair; update one provider to backfill") +} + +// routerConfig mirrors the on-wire shape llm_router accepts. Kept +// private so the synthesiser owns the contract; the proxy-side factory +// JSON-decodes the same shape. +type routerConfig struct { + Providers []routerProviderRoute `json:"providers"` +} + +type routerProviderRoute struct { + ID string `json:"id"` + Vendor string `json:"vendor,omitempty"` + Models []string `json:"models"` + UpstreamScheme string `json:"upstream_scheme"` + UpstreamHost string `json:"upstream_host"` + UpstreamPath string `json:"upstream_path,omitempty"` + AuthHeaderName string `json:"auth_header_name"` + AuthHeaderValue string `json:"auth_header_value"` + AllowedGroupIDs []string `json:"allowed_group_ids,omitempty"` + // Vertex marks a Google Vertex AI provider, whose requests carry the + // model in the URL path. The router selects it by path, bypassing the + // model/vendor table. + Vertex bool `json:"vertex,omitempty"` + // Bedrock marks an AWS Bedrock provider, whose requests carry the model in + // the URL path (/model/{id}/{action}). The router selects it by path, + // bypassing the model/vendor table; auth is a static bearer token. + Bedrock bool `json:"bedrock,omitempty"` + // GCPServiceAccountKeyB64 carries a base64-encoded GCP service-account + // JSON key (from a "keyfile::" api_key). When set, the proxy mints + // + refreshes the OAuth token at request time instead of injecting a static + // AuthHeaderValue. + GCPServiceAccountKeyB64 string `json:"gcp_sa_key_b64,omitempty"` +} + +// indexProviderGroups walks the enabled policies and returns, per +// provider id, the sorted union of source group ids across every +// policy that authorises the provider. Providers with no authorising +// policy are absent from the map. The router consumes this to filter +// candidate routes by the caller's group memberships before the +// path-prefix tiebreak runs. +func indexProviderGroups(policies []*types.Policy) map[string][]string { + sets := make(map[string]map[string]struct{}) + for _, policy := range policies { + if policy == nil { + continue + } + for _, providerID := range policy.DestinationProviderIDs { + if providerID == "" { + continue + } + set, ok := sets[providerID] + if !ok { + set = make(map[string]struct{}) + sets[providerID] = set + } + for _, group := range policy.SourceGroups { + if group != "" { + set[group] = struct{}{} + } + } + } + } + out := make(map[string][]string, len(sets)) + for providerID, set := range sets { + groups := make([]string, 0, len(set)) + for g := range set { + groups = append(groups, g) + } + sort.Strings(groups) + out[providerID] = groups + } + return out +} + +// buildRouterConfigJSON denormalises the account's enabled providers +// into the router middleware's first-match-wins routing table. +// Providers are listed in created_at order so the table is +// deterministic and stable across synth cycles. +// +// AllowedGroupIDs is the union of source group ids across every enabled +// policy that authorises the provider. The router uses it as a hard +// filter — a route whose AllowedGroupIDs has no intersection with the +// caller's user groups is removed from the candidate list before the +// path-prefix tiebreak. Providers no enabled policy authorises +// (orphans) are intentionally OMITTED so the router never observes a +// route with an empty ACL. +func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][]string) ([]byte, error) { + cfg := routerConfig{Providers: make([]routerProviderRoute, 0, len(providers))} + for _, p := range providers { + groups, hasPolicy := groupIndex[p.ID] + if !hasPolicy { + // Orphan: skip. No enabled policy authorises this + // provider, so it must not be reachable. + continue + } + scheme, host, path, err := parseUpstreamHost(p.UpstreamURL) + if err != nil { + return nil, fmt.Errorf("router config for provider %s: %w", p.ID, err) + } + headerName, headerValue, gcpSAKeyB64, err := providerAuthHeader(p) + if err != nil { + return nil, err + } + cfg.Providers = append(cfg.Providers, routerProviderRoute{ + ID: p.ID, + Vendor: providerVendor(p), + Models: providerModelIDs(p), + UpstreamScheme: scheme, + UpstreamHost: host, + UpstreamPath: path, + AuthHeaderName: headerName, + AuthHeaderValue: headerValue, + AllowedGroupIDs: groups, + Vertex: catalog.IsVertexPathStyle(p.ProviderID), + Bedrock: catalog.IsBedrockPathStyle(p.ProviderID), + GCPServiceAccountKeyB64: gcpSAKeyB64, + }) + } + out, err := json.Marshal(cfg) + if err != nil { + return nil, fmt.Errorf("marshal llm_router middleware config: %w", err) + } + return out, nil +} + +// providerVendor returns the parser surface ("openai", "anthropic", …) +// the provider speaks, sourced from its catalog entry's ParserID. The +// router uses it to keep a request the parser tagged with a vendor on a +// route of the same vendor — so e.g. an Anthropic /v1/messages call is +// never sent to an OpenAI-compatible gateway that also claims the model. +// Empty when the catalog entry is unknown or declares no parser surface; +// the router then falls back to model / path routing. +func providerVendor(p *types.Provider) string { + entry, ok := catalog.Lookup(p.ProviderID) + if !ok { + return "" + } + return entry.ParserID +} + +// providerModelIDs returns the model identifiers exposed by the +// provider, deduplicated and in the operator's declared order. Empty +// slice when no models are configured — the router treats that as +// "claim every model" so gateway-style providers (LiteLLM, custom +// OpenAI-compatible endpoints) work without the operator enumerating +// the upstream's full model catalog in NetBird. +func providerModelIDs(p *types.Provider) []string { + if len(p.Models) == 0 { + return []string{} + } + seen := make(map[string]struct{}, len(p.Models)) + out := make([]string, 0, len(p.Models)) + for _, m := range p.Models { + if m.ID == "" { + continue + } + if _, dup := seen[m.ID]; dup { + continue + } + seen[m.ID] = struct{}{} + out = append(out, m.ID) + } + return out +} + +// identityInjectConfig mirrors the on-wire shape llm_identity_inject +// accepts. +type identityInjectConfig struct { + Providers []identityInjectProvider `json:"providers"` +} + +// identityInjectProvider carries one provider's injection rule. +// Identity-stamping uses one of HeaderPair / JSONMetadata (mutually +// exclusive). ExtraHeaders is independent — a list of extra +// per-provider routing/config headers (catalog-declared, value lives +// on the provider record) the middleware stamps with anti-spoof +// (Remove + Add) on every matching request. +type identityInjectProvider struct { + ProviderID string `json:"provider_id"` + HeaderPair *identityInjectHeaderPair `json:"header_pair,omitempty"` + JSONMetadata *identityInjectJSONMetadata `json:"json_metadata,omitempty"` + ExtraHeaders []identityInjectExtraHeader `json:"extra_headers,omitempty"` +} + +type identityInjectExtraHeader struct { + Name string `json:"name"` + Value string `json:"value"` +} + +type identityInjectHeaderPair struct { + EndUserIDHeader string `json:"end_user_id_header,omitempty"` + TagsHeader string `json:"tags_header,omitempty"` + TagsInBody bool `json:"tags_in_body,omitempty"` + EndUserIDInBody bool `json:"end_user_id_in_body,omitempty"` +} + +type identityInjectJSONMetadata struct { + Header string `json:"header"` + UserKey string `json:"user_key,omitempty"` + GroupsKey string `json:"groups_key,omitempty"` + MaxValueLength int `json:"max_value_length,omitempty"` +} + +// buildIdentityInjectConfigJSON walks the enabled providers and emits +// one entry per provider whose catalog entry declares an +// IdentityInjection block. The middleware no-ops for any provider not +// in this list, so the chain is safe to ship to all targets even when +// no identity-stamping provider is configured. +// +// The caller passes groupIndex so we can mirror the synthesiser's own +// "drop orphans" rule — providers no enabled policy authorises don't +// reach the router, so injecting identity for them would never fire. +// We could leave them in for symmetry, but skipping is cheaper and +// clearer. +func buildIdentityInjectConfigJSON(providers []*types.Provider, groupIndex map[string][]string) ([]byte, error) { + cfg := identityInjectConfig{Providers: make([]identityInjectProvider, 0)} + for _, p := range providers { + if _, hasPolicy := groupIndex[p.ID]; !hasPolicy { + continue + } + entry, ok := catalog.Lookup(p.ProviderID) + if !ok { + continue + } + rule := identityInjectProvider{ProviderID: p.ID} + // Identity-stamping shape (one of HeaderPair / JSONMetadata). + // Skip the shape silently when the catalog entry doesn't + // declare one — extras can still apply, see below. + if entry.IdentityInjection != nil { + switch { + case entry.IdentityInjection.HeaderPair != nil: + hp := entry.IdentityInjection.HeaderPair + // For Customizable shapes (Bifrost today) the wire + // header names come from the provider record verbatim; + // the catalog values are placeholder defaults shown by + // the dashboard, not authoritative. Empty operator + // value disables stamping for that dimension — + // applyHeaderPair already no-ops on empty header + // names. The body-inject flags stay catalog-owned + // because Customizable=true today only applies to + // gateways that read identity from headers (the + // flags would be no-ops anyway). + userHeader := hp.EndUserIDHeader + tagsHeader := hp.TagsHeader + if hp.Customizable { + userHeader = p.IdentityHeaderUserID + tagsHeader = p.IdentityHeaderGroups + } + if userHeader != "" || tagsHeader != "" || hp.TagsInBody || hp.EndUserIDInBody { + rule.HeaderPair = &identityInjectHeaderPair{ + EndUserIDHeader: userHeader, + TagsHeader: tagsHeader, + TagsInBody: hp.TagsInBody, + EndUserIDInBody: hp.EndUserIDInBody, + } + } + case entry.IdentityInjection.JSONMetadata != nil: + jm := entry.IdentityInjection.JSONMetadata + // Customizable JSONMetadata reuses the same provider- + // record fields HeaderPair uses — IdentityHeaderUserID + // becomes the JSON key for the user dimension, and + // IdentityHeaderGroups becomes the JSON key for groups. + // Empty operator value is honored as "skip this key"; + // applyJSONMetadata already drops keys with empty + // names. Header itself is catalog-owned (e.g. + // cf-aig-metadata) — operators only override the keys + // inside the JSON, not the wire header that carries it. + userKey := jm.UserKey + groupsKey := jm.GroupsKey + if jm.Customizable { + userKey = p.IdentityHeaderUserID + groupsKey = p.IdentityHeaderGroups + } + if jm.Header != "" { + rule.JSONMetadata = &identityInjectJSONMetadata{ + Header: jm.Header, + UserKey: userKey, + GroupsKey: groupsKey, + MaxValueLength: jm.MaxValueLength, + } + } + } + } + // Extra catalog-declared static headers (e.g. Portkey config + // id). Only emit entries whose value the operator has filled + // in on the provider record; missing/empty values are no-ops. + for _, h := range entry.ExtraHeaders { + if h.Name == "" { + continue + } + v := strings.TrimSpace(p.ExtraValues[h.Name]) + if v == "" { + continue + } + rule.ExtraHeaders = append(rule.ExtraHeaders, identityInjectExtraHeader{ + Name: h.Name, + Value: v, + }) + } + // If this provider would emit nothing, skip it entirely so + // the middleware doesn't carry an inert rule for it. + if rule.HeaderPair == nil && rule.JSONMetadata == nil && len(rule.ExtraHeaders) == 0 { + continue + } + cfg.Providers = append(cfg.Providers, rule) + } + out, err := json.Marshal(cfg) + if err != nil { + return nil, fmt.Errorf("marshal llm_identity_inject middleware config: %w", err) + } + return out, nil +} + +// buildMiddlewareChain assembles the per-target middleware chain that +// implements the Agent Network behaviour at the proxy. Slot order on +// the request leg is the slice order; on the response leg it runs in +// reverse, so cost_meter must come BEFORE llm_response_parser so the +// parser populates token counts before the cost meter reads them. +// +// Authorisation is fused into llm_router: the router carries +// AllowedGroupIDs per provider and filters candidates by the caller's +// user-groups before the path-prefix tiebreak. Per-policy +// enforcement (token / budget caps) lives in llm_limit_check, which +// runs after the router so it can read the resolved provider id; +// llm_limit_record on the response leg posts deltas back to +// management to keep the consumption counters fresh. +// +// llm_identity_inject runs immediately after the router so the +// resolved provider id is available; it stamps NetBird identity onto +// requests bound for gateways like LiteLLM that key budgets and +// attribution off request headers. CanMutate is required so its +// HeadersAdd / HeadersRemove pass the framework's mutation gate. +func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON []byte, redactPii, capturePromptContent bool) []rpservice.MiddlewareConfig { + // Both parsers receive an explicit capture flag derived from the account's + // enable_prompt_collection toggle; nil/unset would default to the legacy + // "always emit" behavior in the middleware, which is precisely what we + // must suppress when the operator hasn't opted in. The flag is duplicated + // across both parsers under distinct field names (capture_prompt / + // capture_completion) to keep each parser's config independently + // auditable. + requestParserCfg := buildParserConfigJSON("capture_prompt", redactPii, capturePromptContent) + responseParserCfg := buildParserConfigJSON("capture_completion", redactPii, capturePromptContent) + return []rpservice.MiddlewareConfig{ + { + ID: middlewareIDLLMRequestParser, + Enabled: true, + Slot: rpservice.MiddlewareSlotOnRequest, + ConfigJSON: requestParserCfg, + }, + { + ID: middlewareIDLLMRouter, + Enabled: true, + Slot: rpservice.MiddlewareSlotOnRequest, + ConfigJSON: routerCfgJSON, + // llm_router rewrites the request's headers (strip + // client auth + inject provider auth) and the upstream + // target via Mutations.RewriteUpstream. Both gated on + // CanMutate; without this flag the chain framework + // drops every mutation and the reverse proxy dials the + // placeholder noop.invalid host (502). + CanMutate: true, + }, + { + // llm_limit_check runs after the router so it knows the + // resolved provider id, but before identity_inject so a + // cap-deny doesn't pay the cost of stamping headers + // we'll never use. + ID: middlewareIDLLMLimitCheck, + Enabled: true, + Slot: rpservice.MiddlewareSlotOnRequest, + ConfigJSON: []byte("{}"), + }, + { + ID: middlewareIDLLMIdentityInject, + Enabled: true, + Slot: rpservice.MiddlewareSlotOnRequest, + ConfigJSON: identityInjectJSON, + // CanMutate is required so HeadersAdd / HeadersRemove + // emitted to stamp NetBird identity onto the upstream + // request actually land — without it the framework + // drops every header mutation. + CanMutate: true, + }, + { + ID: middlewareIDLLMGuardrail, + Enabled: true, + Slot: rpservice.MiddlewareSlotOnRequest, + ConfigJSON: guardrailJSON, + }, + { + // Response slot runs in reverse slice order at runtime: + // limit_record sits FIRST in the response section so it + // runs LAST, after llm_response_parser stamped tokens + // and cost_meter computed cost — both of which the + // recorder reads from the metadata bag. + ID: middlewareIDLLMLimitRecord, + Enabled: true, + Slot: rpservice.MiddlewareSlotOnResponse, + ConfigJSON: []byte("{}"), + }, + { + ID: middlewareIDCostMeter, + Enabled: true, + Slot: rpservice.MiddlewareSlotOnResponse, + ConfigJSON: []byte("{}"), + }, + { + ID: middlewareIDLLMResponseParser, + Enabled: true, + Slot: rpservice.MiddlewareSlotOnResponse, + ConfigJSON: responseParserCfg, + }, + } +} + +// guardrailConfig is the JSON shape the proxy-side llm_guardrail +// middleware expects. Mirrors the proxy registration documented in +// the management→proxy contract. +type guardrailConfig struct { + ModelAllowlist []string `json:"model_allowlist,omitempty"` + PromptCapture guardrailPromptCapture `json:"prompt_capture"` +} + +type guardrailPromptCapture struct { + Enabled bool `json:"enabled"` + RedactPii bool `json:"redact_pii"` +} + +// buildParserConfigJSON assembles the request- or response-parser config JSON. +// captureField names the parser-specific gate (capture_prompt for the request +// parser, capture_completion for the response parser); both are sourced from +// settings.EnablePromptCollection. redact_pii is only meaningful when capture +// is on (no content → nothing to redact) but we forward it verbatim so the +// proxy-side parser stays the only place that interprets the combination. +func buildParserConfigJSON(captureField string, redactPii, capture bool) []byte { + payload := map[string]any{ + captureField: capture, + } + if redactPii { + payload["redact_pii"] = true + } + out, err := json.Marshal(payload) + if err != nil { + // json.Marshal on a map[string]any of bools cannot fail; if it + // somehow does, ship the static minimal config so synth keeps + // working instead of panicking. + return []byte(`{}`) + } + return out +} + +// applyAccountCollectionControls folds the account-level collection master +// switches into the merged guardrail set. Prompt capture enablement is sourced +// SOLELY from the account toggle — the account-network setting is the master +// enable, and policies don't need to attach a capture-enabled guardrail to opt +// in. PII redaction is safe-additive: it applies when either the account or a +// policy guardrail enables it (OR). +func applyAccountCollectionControls(merged *MergedGuardrails, settings *types.Settings) { + if settings == nil { + return + } + merged.PromptCapture.Enabled = settings.EnablePromptCollection + merged.PromptCapture.RedactPii = settings.RedactPii || merged.PromptCapture.RedactPii +} + +func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) { + cfg := guardrailConfig{ + ModelAllowlist: merged.ModelAllowlist, + PromptCapture: guardrailPromptCapture{ + Enabled: merged.PromptCapture.Enabled, + RedactPii: merged.PromptCapture.RedactPii, + }, + } + out, err := json.Marshal(cfg) + if err != nil { + return nil, fmt.Errorf("marshal guardrail middleware config: %w", err) + } + return out, nil +} + +// buildAccountService composes the per-account gateway Service. The +// target carries the noop placeholder URL — the router middleware +// rewrites every request to the matched provider's upstream before the +// proxy dials — alongside the full middleware chain and capture caps. +func buildAccountService( + accountID string, + settings *types.Settings, + enabledPolicies []*types.Policy, + middlewares []rpservice.MiddlewareConfig, + sessionPriv, sessionPub string, +) *rpservice.Service { + cluster := settings.Cluster + domain := settings.Endpoint() + serviceID := SynthesizedServiceIDPrefix + accountID + + return &rpservice.Service{ + ID: serviceID, + AccountID: accountID, + Name: "agent-network-" + accountID, + Domain: domain, + ProxyCluster: cluster, + Mode: rpservice.ModeHTTP, + Enabled: true, + Private: true, + // AccessGroups gates tunnel-peer access (ValidateTunnelPeer) to the + // synthesised agent-network endpoint. Agents reach the gateway over + // the WireGuard tunnel and are authorised by their peer→user group + // membership — the union of every enabled policy's source groups. + AccessGroups: unionSourceGroups(enabledPolicies), + PassHostHeader: false, + RewriteRedirects: false, + SessionPrivateKey: sessionPriv, + SessionPublicKey: sessionPub, + Targets: []*rpservice.Target{ + { + AccountID: accountID, + ServiceID: serviceID, + TargetType: rpservice.TargetTypeCluster, + TargetId: cluster, + Host: noopUpstreamHost, + Port: noopUpstreamPort, + Protocol: noopUpstreamScheme, + Enabled: true, + Options: rpservice.TargetOptions{ + DirectUpstream: true, + AgentNetwork: true, + DisableAccessLog: !settings.EnableLogCollection, + Middlewares: middlewares, + CaptureMaxRequestBytes: agentNetworkRequestCaptureBytes, + CaptureMaxResponseBytes: agentNetworkResponseCaptureBytes, + CaptureContentTypes: append([]string(nil), agentNetworkCaptureContentTypes...), + }, + }, + }, + } +} + +const ( + noopUpstreamScheme = "https" + noopUpstreamHost = "noop.invalid" + noopUpstreamPort = uint16(443) +) + +// providerAuthHeader builds the upstream auth header pair for a +// provider from its catalog entry. The catalog declares which header +// name and template a provider's API expects; the synthesiser +// substitutes the provider's decrypted API key into the template and +// returns the (name, value) pair the router middleware injects after +// stripping the inbound vendor auth headers. +func providerAuthHeader(p *types.Provider) (name, value, gcpSAKeyB64 string, err error) { + entry, ok := catalog.Lookup(p.ProviderID) + if !ok { + return "", "", "", fmt.Errorf("provider %s references unknown catalog id %q", p.ID, p.ProviderID) + } + if entry.AuthHeaderName == "" || entry.AuthHeaderTemplate == "" { + return "", "", "", fmt.Errorf("catalog entry %q has no auth header configured", p.ProviderID) + } + if p.APIKey == "" { + return "", "", "", fmt.Errorf("provider %s has no api key", p.ID) + } + // A "keyfile::" api_key is a GCP service-account key, not a + // static bearer. The proxy mints + refreshes a short-lived OAuth token from + // it at request time, so carry the key material on the route and emit no + // static value. + if rest, isKeyfile := strings.CutPrefix(p.APIKey, gcpKeyfilePrefix); isKeyfile { + return entry.AuthHeaderName, "", strings.TrimSpace(rest), nil + } + value = strings.ReplaceAll(entry.AuthHeaderTemplate, apiKeyPlaceholder, p.APIKey) + return entry.AuthHeaderName, value, "", nil +} + +// parseUpstreamHost splits provider.UpstreamURL into (scheme, host, path) +// where host carries an explicit ":port" suffix when the URL set one +// and path is the URL's path component normalised by stripping a +// trailing slash. The router uses path to disambiguate providers that +// claim the same model. Used by the router config so the rewrite +// carries an authority the reverse proxy can dial verbatim. +func parseUpstreamHost(raw string) (scheme, host, path string, err error) { + parsed, perr := url.Parse(strings.TrimSpace(raw)) + if perr != nil { + return "", "", "", fmt.Errorf("parse upstream_url %q: %w", raw, perr) + } + switch strings.ToLower(parsed.Scheme) { + case "http": + scheme = "http" + case "https": + scheme = "https" + default: + return "", "", "", fmt.Errorf("upstream_url scheme must be http or https, got %q", parsed.Scheme) + } + hostname := parsed.Hostname() + if hostname == "" { + return "", "", "", fmt.Errorf("upstream_url %q has no host", raw) + } + if port := parsed.Port(); port != "" { + host = hostname + ":" + port + } else { + host = hostname + } + path = strings.TrimRight(parsed.Path, "/") + return scheme, host, path, nil +} + +// unionSourceGroups deduplicates source-group IDs across the policies +// pointing at any provider, in deterministic order. +func unionSourceGroups(policies []*types.Policy) []string { + seen := make(map[string]struct{}) + for _, policy := range policies { + for _, group := range policy.SourceGroups { + if group == "" { + continue + } + seen[group] = struct{}{} + } + } + out := make([]string, 0, len(seen)) + for group := range seen { + out = append(out, group) + } + sort.Strings(out) + return out +} + +// MergedGuardrails is the JSON shape passed to the proxy via the +// guardrail middleware's config_json. Mirrors the proxy-side +// expectations and is intentionally distinct from +// types.GuardrailChecks so we can evolve either side independently. +type MergedGuardrails struct { + ModelAllowlist []string `json:"model_allowlist,omitempty"` + TokenLimits MergedTokenLimits `json:"token_limits"` + Budget MergedBudget `json:"budget"` + PromptCapture MergedPromptCapture `json:"prompt_capture"` + Retention MergedRetention `json:"retention"` +} + +type MergedTokenLimits struct { + Hourly *MergedTokenWindow `json:"hourly,omitempty"` + Daily *MergedTokenWindow `json:"daily,omitempty"` + Monthly *MergedTokenWindow `json:"monthly,omitempty"` +} + +type MergedTokenWindow struct { + MaxInputTokens int `json:"max_input_tokens,omitempty"` + MaxOutputTokens int `json:"max_output_tokens,omitempty"` +} + +type MergedBudget struct { + Hourly *MergedBudgetWindow `json:"hourly,omitempty"` + Daily *MergedBudgetWindow `json:"daily,omitempty"` + Monthly *MergedBudgetWindow `json:"monthly,omitempty"` +} + +type MergedBudgetWindow struct { + SoftCapUSD float64 `json:"soft_cap_usd,omitempty"` + HardCapUSD float64 `json:"hard_cap_usd,omitempty"` +} + +type MergedPromptCapture struct { + Enabled bool `json:"enabled"` + RedactPii bool `json:"redact_pii"` +} + +type MergedRetention struct { + Enabled bool `json:"enabled"` + Days int `json:"days"` +} + +// mergeGuardrails computes the effective guardrail spec applied at the +// proxy, given the referencing policies and the account's guardrail +// catalogue. Policy enabled-ness is the caller's responsibility — only +// enabled policies should be passed in. +// +// Merge rules: +// - Model allowlist: union of allowlists across policies that enable it. +// - Token / Budget: most-restrictive (min of non-zero caps) per window. +// - Prompt capture: enabled if any policy enables it; redact_pii sticks +// if any enabling policy turns it on. +// - Retention: enabled if any enables it; smallest non-zero days wins. +func mergeGuardrails(policies []*types.Policy, byID map[string]*types.Guardrail) MergedGuardrails { + merged := MergedGuardrails{} + allowlist := make(map[string]struct{}) + allowlistEnabled := false + + for _, policy := range policies { + for _, gID := range policy.GuardrailIDs { + g, ok := byID[gID] + if !ok || g == nil { + continue + } + if g.Checks.ModelAllowlist.Enabled { + allowlistEnabled = true + for _, m := range g.Checks.ModelAllowlist.Models { + if m != "" { + allowlist[m] = struct{}{} + } + } + } + if g.Checks.PromptCapture.Enabled { + merged.PromptCapture.Enabled = true + if g.Checks.PromptCapture.RedactPii { + merged.PromptCapture.RedactPii = true + } + } + // TokenLimits, Budget, and Retention have moved off + // guardrails. Token and budget caps now live on the + // Policy itself (Policy.Limits); retention will move to + // account-level Settings. The proxy still consumes + // MergedTokenLimits/MergedBudget/MergedRetention from + // this struct for backwards compatibility, but they + // remain at zero/disabled until the new Policy.Limits + // merge is wired in by the enforcement track. + } + } + + if allowlistEnabled { + merged.ModelAllowlist = make([]string, 0, len(allowlist)) + for m := range allowlist { + merged.ModelAllowlist = append(merged.ModelAllowlist, m) + } + sort.Strings(merged.ModelAllowlist) + } + return merged +} diff --git a/management/server/agentnetwork/synthesizer_guardrail_realstore_test.go b/management/server/agentnetwork/synthesizer_guardrail_realstore_test.go new file mode 100644 index 000000000..7c7bb8bcf --- /dev/null +++ b/management/server/agentnetwork/synthesizer_guardrail_realstore_test.go @@ -0,0 +1,178 @@ +package agentnetwork + +import ( + "context" + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/store" +) + +// decodeServiceGuardrailConfig pulls the llm_guardrail middleware config off the +// synthesised service's single target. +func decodeServiceGuardrailConfig(t *testing.T, svc *rpservice.Service) guardrailConfig { + t.Helper() + require.NotEmpty(t, svc.Targets, "synth service must carry a target") + for _, mw := range svc.Targets[0].Options.Middlewares { + if mw.ID == middlewareIDLLMGuardrail { + var cfg guardrailConfig + require.NoError(t, json.Unmarshal(mw.ConfigJSON, &cfg), "guardrail config must decode") + return cfg + } + } + t.Fatal("llm_guardrail middleware not present on synthesised service") + return guardrailConfig{} +} + +// decodeMiddlewareRawConfig returns the raw ConfigJSON bytes for the named +// middleware on the synth service's target, or fails the test. +func decodeMiddlewareRawConfig(t *testing.T, svc *rpservice.Service, id string) []byte { + t.Helper() + require.NotEmpty(t, svc.Targets, "synth service must carry a target") + for _, mw := range svc.Targets[0].Options.Middlewares { + if mw.ID == id { + return mw.ConfigJSON + } + } + t.Fatalf("middleware %q not present on synthesised service", id) + return nil +} + +// saveGuardrailAndPolicy persists a guardrail with prompt capture + redact + a +// model allowlist, referenced by one enabled policy. Shared by the GC-3 tests. +func saveGuardrailAndPolicy(t *testing.T, ctx context.Context, s store.Store, provider *types.Provider) { + t.Helper() + guardrail := &types.Guardrail{ + ID: "ainguard-1", + AccountID: testAccountID, + Name: "strict", + Checks: types.GuardrailChecks{ + ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true, Models: []string{"gpt-5.4"}}, + PromptCapture: types.GuardrailPromptCapture{Enabled: true, RedactPii: true}, + }, + } + require.NoError(t, s.SaveAgentNetworkGuardrail(ctx, guardrail)) + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, newSynthTestPolicy(provider.ID, "grp-eng", guardrail.ID))) +} + +// TestSynthesizeServices_RealStore_PromptCaptureAccountIsSoleControl is the +// GC-3 contract: the account master switch (EnablePromptCollection) is the +// SOLE control for capture enablement. Policy-level guardrail prompt_capture is +// ignored for enablement — operators don't need to attach a capture guardrail +// to a policy just to turn capture on for the account. Off by default. +func TestSynthesizeServices_RealStore_PromptCaptureAccountIsSoleControl(t *testing.T) { + ctx := context.Background() + s, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err, "real sqlite test store must come up") + defer cleanup() + + // Account collection master switch OFF (default). + require.NoError(t, s.SaveAgentNetworkSettings(ctx, newSynthTestSettings())) + saveGuardrailAndPolicy(t, ctx, s, newSynthTestProvider()) + + services, err := SynthesizeServices(ctx, s, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + cfg := decodeServiceGuardrailConfig(t, services[0]) + assert.Equal(t, []string{"gpt-5.4"}, cfg.ModelAllowlist, + "model allowlist is a pure policy guardrail and must always reach the config") + assert.False(t, cfg.PromptCapture.Enabled, + "prompt capture must be off when the account toggle is off, even with a capture-enabled guardrail") +} + +// TestSynthesizeServices_RealStore_PromptCaptureFlowsWhenAccountOptsIn proves +// the account toggle is sufficient on its own — even with NO guardrail +// attached to the policy, capture fires when the account opts in. Redact is +// the OR of account + guardrail. +func TestSynthesizeServices_RealStore_PromptCaptureFlowsWhenAccountOptsIn(t *testing.T) { + ctx := context.Background() + s, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err, "real sqlite test store must come up") + defer cleanup() + + settings := newSynthTestSettings() + settings.EnablePromptCollection = true + require.NoError(t, s.SaveAgentNetworkSettings(ctx, settings)) + + // Save a provider and a policy with NO guardrails attached — proves the + // account toggle is sufficient on its own. + provider := newSynthTestProvider() + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, newSynthTestPolicy(provider.ID, "grp-eng", ""))) + + services, err := SynthesizeServices(ctx, s, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + cfg := decodeServiceGuardrailConfig(t, services[0]) + assert.True(t, cfg.PromptCapture.Enabled, + "account toggle alone must enable capture; no guardrail attachment required") +} + +// TestSynthesizeServices_RealStore_AccountRedactWithoutGuardrailRedact proves +// the redact OR-merge from the account side: account RedactPii on, guardrail +// redact off, capture on at both levels. +func TestSynthesizeServices_RealStore_AccountRedactWithoutGuardrailRedact(t *testing.T) { + ctx := context.Background() + s, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err, "real sqlite test store must come up") + defer cleanup() + + settings := newSynthTestSettings() + settings.EnablePromptCollection = true + settings.RedactPii = true + require.NoError(t, s.SaveAgentNetworkSettings(ctx, settings)) + + provider := newSynthTestProvider() + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + guardrail := &types.Guardrail{ + ID: "ainguard-noredact", + AccountID: testAccountID, + Name: "capture-only", + Checks: types.GuardrailChecks{ + PromptCapture: types.GuardrailPromptCapture{Enabled: true, RedactPii: false}, + }, + } + require.NoError(t, s.SaveAgentNetworkGuardrail(ctx, guardrail)) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, newSynthTestPolicy(provider.ID, "grp-eng", guardrail.ID))) + + services, err := SynthesizeServices(ctx, s, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + cfg := decodeServiceGuardrailConfig(t, services[0]) + assert.True(t, cfg.PromptCapture.Enabled, "capture on (account + guardrail)") + assert.True(t, cfg.PromptCapture.RedactPii, "account RedactPii must apply even when the guardrail leaves it off (OR)") +} + +// TestSynthesizeServices_RealStore_NoGuardrail_CaptureOff pins the default: +// with no guardrail referenced, the synth service's guardrail config has prompt +// capture disabled and an empty allowlist. This is the "off by default" baseline +// the account switch must preserve. +func TestSynthesizeServices_RealStore_NoGuardrail_CaptureOff(t *testing.T) { + ctx := context.Background() + s, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err, "real sqlite test store must come up") + defer cleanup() + + require.NoError(t, s.SaveAgentNetworkSettings(ctx, newSynthTestSettings())) + provider := newSynthTestProvider() + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, newSynthTestPolicy(provider.ID, "grp-eng", ""))) + + services, err := SynthesizeServices(ctx, s, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1, "exactly one synth service expected") + + cfg := decodeServiceGuardrailConfig(t, services[0]) + assert.Empty(t, cfg.ModelAllowlist, "no guardrail → no allowlist") + assert.False(t, cfg.PromptCapture.Enabled, "no guardrail → prompt capture off by default") + assert.False(t, cfg.PromptCapture.RedactPii, "no guardrail → redact off by default") +} diff --git a/management/server/agentnetwork/synthesizer_log_collection_realstore_test.go b/management/server/agentnetwork/synthesizer_log_collection_realstore_test.go new file mode 100644 index 000000000..9aa2a0abe --- /dev/null +++ b/management/server/agentnetwork/synthesizer_log_collection_realstore_test.go @@ -0,0 +1,70 @@ +package agentnetwork + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy" + rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" + "github.com/netbirdio/netbird/management/server/store" +) + +// TestSynthesizeServices_RealStore_LogCollectionOff_SuppressesAccessLog drives the +// happy default: account settings ship with EnableLogCollection=false, so the +// synthesised target opts out of access-log emission (DisableAccessLog=true) and +// the proto mapping the proxy receives reflects that. +func TestSynthesizeServices_RealStore_LogCollectionOff_SuppressesAccessLog(t *testing.T) { + ctx := context.Background() + s, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err, "real sqlite test store must come up") + defer cleanup() + + require.NoError(t, s.SaveAgentNetworkSettings(ctx, newSynthTestSettings())) + provider := newSynthTestProvider() + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, newSynthTestPolicy(provider.ID, "grp-eng", ""))) + + services, err := SynthesizeServices(ctx, s, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1, "exactly one synth service expected") + require.NotEmpty(t, services[0].Targets, "synth service must carry a target") + assert.True(t, services[0].Targets[0].Options.DisableAccessLog, + "EnableLogCollection=false (default) must produce DisableAccessLog=true on the synth target") + + mapping := services[0].ToProtoMapping(rpservice.Update, "", rpproxy.OIDCValidationConfig{}) + require.NotEmpty(t, mapping.GetPath(), "proto mapping must carry a path") + assert.True(t, mapping.GetPath()[0].GetOptions().GetDisableAccessLog(), + "proto mapping must propagate DisableAccessLog=true so the proxy suppresses access-log emission") +} + +// TestSynthesizeServices_RealStore_LogCollectionOn_PermitsAccessLog asserts the +// inverse: once the account opts in, the synth target leaves DisableAccessLog +// at its default false and the proto wire stays unset. +func TestSynthesizeServices_RealStore_LogCollectionOn_PermitsAccessLog(t *testing.T) { + ctx := context.Background() + s, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err, "real sqlite test store must come up") + defer cleanup() + + settings := newSynthTestSettings() + settings.EnableLogCollection = true + require.NoError(t, s.SaveAgentNetworkSettings(ctx, settings)) + provider := newSynthTestProvider() + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, newSynthTestPolicy(provider.ID, "grp-eng", ""))) + + services, err := SynthesizeServices(ctx, s, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1, "exactly one synth service expected") + require.NotEmpty(t, services[0].Targets, "synth service must carry a target") + assert.False(t, services[0].Targets[0].Options.DisableAccessLog, + "EnableLogCollection=true must leave DisableAccessLog=false on the synth target") + + mapping := services[0].ToProtoMapping(rpservice.Update, "", rpproxy.OIDCValidationConfig{}) + require.NotEmpty(t, mapping.GetPath(), "proto mapping must carry a path") + assert.False(t, mapping.GetPath()[0].GetOptions().GetDisableAccessLog(), + "proto mapping must propagate DisableAccessLog=false so access-log emission stays on") +} diff --git a/management/server/agentnetwork/synthesizer_parser_redact_realstore_test.go b/management/server/agentnetwork/synthesizer_parser_redact_realstore_test.go new file mode 100644 index 000000000..5f42c3b39 --- /dev/null +++ b/management/server/agentnetwork/synthesizer_parser_redact_realstore_test.go @@ -0,0 +1,145 @@ +package agentnetwork + +import ( + "context" + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" + "github.com/netbirdio/netbird/management/server/store" +) + +// parserRedactConfig mirrors the on-wire shape of the redact + capture knobs +// that both llm_request_parser and llm_response_parser unmarshal. We don't +// import the proxy-side packages from a management test (cross-module), so we +// decode the JSON directly and assert on the fields that are part of the +// synth contract. +type parserRedactConfig struct { + RedactPii bool `json:"redact_pii,omitempty"` + CapturePrompt *bool `json:"capture_prompt,omitempty"` // present only on the request parser + CaptureCompletion *bool `json:"capture_completion,omitempty"` // present only on the response parser +} + +// TestSynthesizeServices_RealStore_ParserConfigsCarryRedactPii is the +// management-side contract test for the request/response parser redaction +// wiring. When settings.RedactPii is true, the synthesised middleware chain +// MUST stamp redact_pii=true on both llm_request_parser and llm_response_parser +// configs — otherwise the parsers ship raw prompts / completions to the +// access log even though the account has opted in. This is exactly the live +// leak path that motivated the parser-side redaction in the first place. +func TestSynthesizeServices_RealStore_ParserConfigsCarryRedactPii(t *testing.T) { + ctx := context.Background() + s, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err, "real sqlite test store must come up") + defer cleanup() + + settings := newSynthTestSettings() + settings.RedactPii = true + settings.EnablePromptCollection = true + require.NoError(t, s.SaveAgentNetworkSettings(ctx, settings)) + + provider := newSynthTestProvider() + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, newSynthTestPolicy(provider.ID, "grp-eng", ""))) + + services, err := SynthesizeServices(ctx, s, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1, "exactly one synth service expected") + + for _, parserID := range []string{middlewareIDLLMRequestParser, middlewareIDLLMResponseParser} { + raw := decodeMiddlewareRawConfig(t, services[0], parserID) + var cfg parserRedactConfig + require.NoError(t, json.Unmarshal(raw, &cfg), "%s config must be valid JSON", parserID) + assert.True(t, cfg.RedactPii, "%s config must carry redact_pii=true when settings.RedactPii is on (otherwise the parser ships raw prompts/completions to the access log)", parserID) + } + // The capture flag is set explicitly to enable_prompt_collection on each + // parser. With it on here, both must allow emission. + reqCfg := decodeParserConfig(t, services[0], middlewareIDLLMRequestParser) + require.NotNil(t, reqCfg.CapturePrompt, "request parser must carry an explicit capture_prompt") + assert.True(t, *reqCfg.CapturePrompt, "capture_prompt=true when EnablePromptCollection=true") + respCfg := decodeParserConfig(t, services[0], middlewareIDLLMResponseParser) + require.NotNil(t, respCfg.CaptureCompletion, "response parser must carry an explicit capture_completion") + assert.True(t, *respCfg.CaptureCompletion, "capture_completion=true when EnablePromptCollection=true") +} + +// decodeParserConfig is a small helper around decodeMiddlewareRawConfig that +// also unmarshals into parserRedactConfig. +func decodeParserConfig(t *testing.T, svc *rpservice.Service, parserID string) parserRedactConfig { + t.Helper() + raw := decodeMiddlewareRawConfig(t, svc, parserID) + var cfg parserRedactConfig + require.NoError(t, json.Unmarshal(raw, &cfg), "%s config must be valid JSON", parserID) + return cfg +} + +// TestSynthesizeServices_RealStore_ParserConfigsSuppressCaptureWhenLogCollectionOnly +// is the contract test for the bug: enable_log_collection=true with +// enable_prompt_collection=false MUST result in capture_prompt=false on the +// request parser AND capture_completion=false on the response parser, so the +// access-log row stays metadata-only (provider, model, tokens, cost) and +// carries NO prompt input nor response output. Without this, operators who +// want billing-style logs end up with raw user prompts and model outputs in +// every access-log entry. +func TestSynthesizeServices_RealStore_ParserConfigsSuppressCaptureWhenLogCollectionOnly(t *testing.T) { + ctx := context.Background() + s, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err, "real sqlite test store must come up") + defer cleanup() + + settings := newSynthTestSettings() + settings.EnableLogCollection = true // operator wants logs ON + settings.EnablePromptCollection = false // but NOT content capture + require.NoError(t, s.SaveAgentNetworkSettings(ctx, settings)) + + provider := newSynthTestProvider() + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, newSynthTestPolicy(provider.ID, "grp-eng", ""))) + + services, err := SynthesizeServices(ctx, s, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + reqCfg := decodeParserConfig(t, services[0], middlewareIDLLMRequestParser) + require.NotNil(t, reqCfg.CapturePrompt, "request parser must carry an explicit capture_prompt gate") + assert.False(t, *reqCfg.CapturePrompt, "capture_prompt MUST be false when EnablePromptCollection is off — otherwise llm.request_prompt_raw leaks user input into the access log") + + respCfg := decodeParserConfig(t, services[0], middlewareIDLLMResponseParser) + require.NotNil(t, respCfg.CaptureCompletion, "response parser must carry an explicit capture_completion gate") + assert.False(t, *respCfg.CaptureCompletion, "capture_completion MUST be false when EnablePromptCollection is off — otherwise llm.response_completion leaks model output into the access log") +} + +// TestSynthesizeServices_RealStore_ParserConfigsOmitRedactPiiWhenOff proves +// the inverse: with the account toggle off, the parser configs stay clean (no +// redact_pii field, which the parsers treat as zero / no redaction). This is +// the operator-opt-out path — the access log keeps raw prompts/completions +// for debugging until the operator opts in. +func TestSynthesizeServices_RealStore_ParserConfigsOmitRedactPiiWhenOff(t *testing.T) { + ctx := context.Background() + s, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err) + defer cleanup() + + // Default settings: RedactPii = false. + require.NoError(t, s.SaveAgentNetworkSettings(ctx, newSynthTestSettings())) + provider := newSynthTestProvider() + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, newSynthTestPolicy(provider.ID, "grp-eng", ""))) + + services, err := SynthesizeServices(ctx, s, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + for _, parserID := range []string{middlewareIDLLMRequestParser, middlewareIDLLMResponseParser} { + raw := decodeMiddlewareRawConfig(t, services[0], parserID) + // Inspect the decoded JSON directly: a struct decode would also pass + // if redact_pii were present-but-false. The contract is that the key + // is omitted entirely while the account toggle is off. + var rawCfg map[string]json.RawMessage + require.NoError(t, json.Unmarshal(raw, &rawCfg), "%s config must be valid JSON", parserID) + assert.NotContains(t, rawCfg, "redact_pii", + "%s config must omit redact_pii entirely while the account toggle is off", parserID) + } +} diff --git a/management/server/agentnetwork/synthesizer_realstore_test.go b/management/server/agentnetwork/synthesizer_realstore_test.go new file mode 100644 index 000000000..1e07c0e81 --- /dev/null +++ b/management/server/agentnetwork/synthesizer_realstore_test.go @@ -0,0 +1,174 @@ +package agentnetwork + +import ( + "context" + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy" + rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" + "github.com/netbirdio/netbird/management/server/account" + "github.com/netbirdio/netbird/management/server/store" + nbtypes "github.com/netbirdio/netbird/management/server/types" + "github.com/netbirdio/netbird/shared/management/proto" +) + +// decodeServiceRouterConfig finds the llm_router middleware on the synthesised +// service's single target and decodes its config — the model→provider routing +// table the proxy authorises against. +func decodeServiceRouterConfig(t *testing.T, svc *rpservice.Service) routerConfig { + t.Helper() + require.NotEmpty(t, svc.Targets, "synth service must carry a target") + for _, mw := range svc.Targets[0].Options.Middlewares { + if mw.ID == middlewareIDLLMRouter { + var cfg routerConfig + require.NoError(t, json.Unmarshal(mw.ConfigJSON, &cfg), "router config must decode") + return cfg + } + } + t.Fatal("llm_router middleware not present on synthesised service") + return routerConfig{} +} + +// decodeMappingRouterConfig is the proto-wire equivalent: it pulls the +// llm_router config off the ProxyMapping the proxy actually receives. +func decodeMappingRouterConfig(t *testing.T, m *proto.ProxyMapping) routerConfig { + t.Helper() + require.NotEmpty(t, m.GetPath(), "mapping must carry a path") + for _, mw := range m.GetPath()[0].GetOptions().GetMiddlewares() { + if mw.GetId() == middlewareIDLLMRouter { + var cfg routerConfig + require.NoError(t, json.Unmarshal(mw.GetConfigJson(), &cfg), "wire router config must decode") + return cfg + } + } + t.Fatal("llm_router middleware not present on proxy mapping") + return routerConfig{} +} + +// TestSynthesizeServices_RealStore_SurvivesStatusToggle drives synthesis through +// a REAL sqlite store (Save → gorm/JSON serialize → reload → decrypt) instead of +// a MockStore, so it exercises the field round-trip that a provider/policy edit +// actually hits. Mock-based tests can't catch a field that dies in persistence; +// this one can. It then performs the exact operation that reproduced the live +// 403 — disable then re-enable the provider — and asserts the re-enabled state +// is fully routable again. +func TestSynthesizeServices_RealStore_SurvivesStatusToggle(t *testing.T) { + ctx := context.Background() + s, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err, "real sqlite test store must come up") + defer cleanup() + + require.NoError(t, s.SaveAgentNetworkSettings(ctx, newSynthTestSettings())) + provider := newSynthTestProvider() + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, newSynthTestPolicy(provider.ID, "grp-eng", ""))) + + assertRoutable := func(t *testing.T, stage string) { + services, err := SynthesizeServices(ctx, s, testAccountID) + require.NoError(t, err, stage) + require.Len(t, services, 1, "%s: exactly one synth service expected", stage) + svc := services[0] + + assert.True(t, svc.Private, "%s: synth service must be Private after store round-trip", stage) + assert.Equal(t, []string{"grp-eng"}, svc.AccessGroups, "%s: AccessGroups must survive the round-trip", stage) + + m := svc.ToProtoMapping(rpservice.Update, "", rpproxy.OIDCValidationConfig{}) + assert.True(t, m.GetPrivate(), "%s: proto mapping Private must be true (proxy gates tunnel-peer auth on it)", stage) + + cfg := decodeServiceRouterConfig(t, svc) + require.Len(t, cfg.Providers, 1, "%s: the enabled+linked provider must appear in the router config", stage) + assert.Equal(t, []string{"gpt-5.4"}, cfg.Providers[0].Models, "%s: provider models must reach the route", stage) + assert.Equal(t, []string{"grp-eng"}, cfg.Providers[0].AllowedGroupIDs, "%s: policy source groups must reach the route", stage) + } + + assertRoutable(t, "initial") + + provider.Enabled = false + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + disabled, err := SynthesizeServices(ctx, s, testAccountID) + require.NoError(t, err, "synthesis must not error with a disabled provider") + for _, svc := range disabled { + assert.Empty(t, decodeServiceRouterConfig(t, svc).Providers, + "a disabled provider must not appear in the router config (otherwise it would route while off)") + } + + provider.Enabled = true + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + assertRoutable(t, "after disable->enable") +} + +// captureController is a proxy.Controller that records the mappings reconcile +// pushes, so the test can inspect the exact wire payload — Private flag and +// router config included. +type captureController struct { + rpproxy.Controller + pushed []*proto.ProxyMapping +} + +func (c *captureController) GetOIDCValidationConfig() rpproxy.OIDCValidationConfig { + return rpproxy.OIDCValidationConfig{} +} + +func (c *captureController) SendServiceUpdateToCluster(_ context.Context, _ string, update *proto.ProxyMapping, _ string) { + c.pushed = append(c.pushed, update) +} + +// noopAccountManager satisfies the reconcile path's accountManager dependency. +type noopAccountManager struct { + account.Manager +} + +func (noopAccountManager) UpdateAccountPeers(context.Context, string, nbtypes.UpdateReason) {} + +// TestReconcile_RealStore_PushesPrivateAfterStatusToggle reproduces the live +// path end-to-end below the gRPC boundary: a real store + the real +// managerImpl.reconcile + a capturing proxy controller. It runs the operation +// that broke in production — provider disable then re-enable — and asserts the +// mapping reconcile pushes to the cluster after re-enable is Private=true and +// carries the routable provider. If reconcile ever pushes private=false (the +// symptom that left UserGroups empty → no_authorised_provider), this fails. +func TestReconcile_RealStore_PushesPrivateAfterStatusToggle(t *testing.T) { + ctx := context.Background() + s, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err) + defer cleanup() + + require.NoError(t, s.SaveAgentNetworkSettings(ctx, newSynthTestSettings())) + provider := newSynthTestProvider() + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + require.NoError(t, s.SaveAgentNetworkPolicy(ctx, newSynthTestPolicy(provider.ID, "grp-eng", ""))) + + ctrl := &captureController{} + m := &managerImpl{ + store: s, + accountManager: noopAccountManager{}, + proxyController: ctrl, + reconcileCache: make(map[string]map[string]*proto.ProxyMapping), + } + + m.reconcile(ctx, testAccountID) // initial, provider enabled + + provider.Enabled = false + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + m.reconcile(ctx, testAccountID) // disabled + + provider.Enabled = true + require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider)) + m.reconcile(ctx, testAccountID) // re-enabled — the reproduction step + + require.NotEmpty(t, ctrl.pushed, "reconcile must push at least one mapping") + last := ctrl.pushed[len(ctrl.pushed)-1] + + assert.Equal(t, newSynthTestSettings().Endpoint(), last.GetDomain(), "synth domain on the wire") + assert.True(t, last.GetPrivate(), + "reconcile-pushed mapping after re-enable MUST be Private=true; a false here is the exact bug — the proxy skips ValidateTunnelPeer, UserGroups stays empty, and llm_router denies no_authorised_provider") + + cfg := decodeMappingRouterConfig(t, last) + require.Len(t, cfg.Providers, 1, "re-enabled provider must be back in the pushed router config") + assert.Equal(t, []string{"gpt-5.4"}, cfg.Providers[0].Models, "model must be routable again after re-enable") + assert.Equal(t, []string{"grp-eng"}, cfg.Providers[0].AllowedGroupIDs, "authorised groups must be present after re-enable") +} diff --git a/management/server/agentnetwork/synthesizer_test.go b/management/server/agentnetwork/synthesizer_test.go new file mode 100644 index 000000000..70e956358 --- /dev/null +++ b/management/server/agentnetwork/synthesizer_test.go @@ -0,0 +1,1098 @@ +package agentnetwork + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/golang/mock/gomock" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/store" + "github.com/netbirdio/netbird/shared/management/status" +) + +const ( + testAccountID = "acct-1" + testCluster = "eu.proxy.netbird.io" + testSubdomain = "violet" + testEndpoint = "violet.eu.proxy.netbird.io" +) + +func newSynthTestSettings() *types.Settings { + return &types.Settings{ + AccountID: testAccountID, + Cluster: testCluster, + Subdomain: testSubdomain, + } +} + +func newSynthTestProvider() *types.Provider { + return &types.Provider{ + ID: "prov-1", + AccountID: testAccountID, + ProviderID: "openai_api", + Name: "OpenAI", + UpstreamURL: "https://api.openai.com", + APIKey: "sk-test-key", + Enabled: true, + Models: []types.ProviderModel{{ID: "gpt-5.4", InputPer1k: 0.0025, OutputPer1k: 0.015}}, + SessionPrivateKey: "test-priv-key", + SessionPublicKey: "test-pub-key", + CreatedAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC), + } +} + +func newSynthTestPolicy(providerID, sourceGroupID, guardrailID string) *types.Policy { + policy := &types.Policy{ + ID: "pol-1", + AccountID: testAccountID, + Name: "engineers", + Enabled: true, + SourceGroups: []string{sourceGroupID}, + DestinationProviderIDs: []string{providerID}, + } + if guardrailID != "" { + policy.GuardrailIDs = []string{guardrailID} + } + return policy +} + +// expectSynthBaseInputs wires the four reads the new synthesiser issues +// in the happy path: settings, providers, policies, guardrails. +func expectSynthBaseInputs(mockStore *store.MockStore, ctx context.Context, settings *types.Settings, providers []*types.Provider, policies []*types.Policy, guardrails []*types.Guardrail) { + if settings == nil { + mockStore.EXPECT(). + GetAgentNetworkSettings(ctx, store.LockingStrengthNone, testAccountID). + Return(nil, status.Errorf(status.NotFound, "agent network settings not found")) + return + } + mockStore.EXPECT(). + GetAgentNetworkSettings(ctx, store.LockingStrengthNone, testAccountID). + Return(settings, nil) + mockStore.EXPECT(). + GetAccountAgentNetworkProviders(ctx, store.LockingStrengthNone, testAccountID). + Return(providers, nil) + if hasEnabled(providers) { + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(ctx, store.LockingStrengthNone, testAccountID). + Return(policies, nil) + if hasEnabledPolicy(policies) { + mockStore.EXPECT(). + GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, testAccountID). + Return(guardrails, nil) + } + } +} + +func hasEnabled(providers []*types.Provider) bool { + for _, p := range providers { + if p != nil && p.Enabled { + return true + } + } + return false +} + +func hasEnabledPolicy(policies []*types.Policy) bool { + for _, p := range policies { + if p != nil && p.Enabled { + return true + } + } + return false +} + +func TestSynthesizeServices_HappyPath(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + openai := newSynthTestProvider() + anthropic := &types.Provider{ + ID: "prov-2", + AccountID: testAccountID, + ProviderID: "anthropic_api", + Name: "Anthropic", + UpstreamURL: "https://api.anthropic.com", + APIKey: "sk-ant-secret", + Enabled: true, + Models: []types.ProviderModel{{ID: "claude-opus-4-7"}}, + SessionPrivateKey: "ant-priv", + SessionPublicKey: "ant-pub", + CreatedAt: time.Date(2026, 2, 1, 0, 0, 0, 0, time.UTC), + } + + policyEng := newSynthTestPolicy(openai.ID, "grp-eng", "") + policyEng.ID = "pol-eng" + policyOps := newSynthTestPolicy(anthropic.ID, "grp-ops", "") + policyOps.ID = "pol-ops" + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{openai, anthropic}, + []*types.Policy{policyEng, policyOps}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err, "synthesis must succeed") + require.Len(t, services, 1, "exactly one service per account") + + svc := services[0] + assert.Equal(t, "agent-net-svc-acct-1", svc.ID, "service id is account-scoped") + assert.Equal(t, testAccountID, svc.AccountID, "service inherits account ID") + assert.Equal(t, testEndpoint, svc.Domain, "domain is settings.Endpoint() (subdomain.cluster)") + assert.Equal(t, testCluster, svc.ProxyCluster, "proxy cluster comes from settings") + assert.Equal(t, rpservice.ModeHTTP, svc.Mode, "synthesised services are HTTP mode") + assert.True(t, svc.Private, "synthesised services are always private") + assert.True(t, svc.Enabled, "synthesised services are enabled when emitted") + assert.Equal(t, []string{"grp-eng", "grp-ops"}, svc.AccessGroups, + "access groups union both policies' source groups (tunnel-peer auth)") + + require.Len(t, svc.Targets, 1, "single cluster target") + target := svc.Targets[0] + assert.Equal(t, rpservice.TargetTypeCluster, target.TargetType, "target type is cluster") + assert.Equal(t, testCluster, target.TargetId, "target id is the cluster address") + assert.Equal(t, "noop.invalid", target.Host, "host is the placeholder; router rewrites at request time") + assert.Equal(t, uint16(443), target.Port, "placeholder port") + assert.Equal(t, "https", target.Protocol, "placeholder scheme") + assert.True(t, target.Options.DirectUpstream, "synth targets imply direct upstream") + assert.True(t, target.Options.AgentNetwork, "synth targets must be flagged as agent_network") + + mws := target.Options.Middlewares + require.Len(t, mws, 8, "eight middlewares: request_parser, router, limit_check, identity_inject, guardrail, limit_record, cost_meter, response_parser") + assert.Equal(t, middlewareIDLLMRequestParser, mws[0].ID, "first middleware is the request parser") + assert.Equal(t, rpservice.MiddlewareSlotOnRequest, mws[0].Slot, "request parser runs on_request") + // Request parser carries the capture_prompt gate sourced from + // settings.EnablePromptCollection. The synth-test settings default + // EnablePromptCollection=false, so capture is off and the access-log row + // will not carry prompt content. + assert.JSONEq(t, `{"capture_prompt":false}`, string(mws[0].ConfigJSON), "request parser config must carry capture_prompt from synth") + + assert.Equal(t, middlewareIDLLMRouter, mws[1].ID, "second middleware is the router") + assert.Equal(t, rpservice.MiddlewareSlotOnRequest, mws[1].Slot, "router runs on_request") + assert.True(t, mws[1].CanMutate, "router must carry CanMutate=true; without it the framework drops the auth-header strip/inject AND the upstream rewrite, leaving the proxy to dial the placeholder noop.invalid") + require.NotEmpty(t, mws[1].ConfigJSON, "router config JSON must be populated") + + var routerCfg routerConfig + require.NoError(t, json.Unmarshal(mws[1].ConfigJSON, &routerCfg), "router config must unmarshal") + require.Len(t, routerCfg.Providers, 2, "both providers must reach the router") + assert.Equal(t, openai.ID, routerCfg.Providers[0].ID, "openai is first by created_at") + assert.Equal(t, "Bearer sk-test-key", routerCfg.Providers[0].AuthHeaderValue, "openai auth header value substitutes the API key") + assert.Equal(t, "Authorization", routerCfg.Providers[0].AuthHeaderName, "openai uses Authorization header") + assert.Equal(t, "https", routerCfg.Providers[0].UpstreamScheme, "openai scheme") + assert.Equal(t, "api.openai.com", routerCfg.Providers[0].UpstreamHost, "openai host") + assert.Equal(t, []string{"grp-eng"}, routerCfg.Providers[0].AllowedGroupIDs, "openai inherits policyEng's source groups") + assert.Equal(t, []string{"gpt-5.4"}, routerCfg.Providers[0].Models, + "the provider's configured model IDs must reach the router route — otherwise the model never matches and llm_router denies model_not_routable") + assert.Equal(t, anthropic.ID, routerCfg.Providers[1].ID, "anthropic follows openai by created_at") + assert.Equal(t, "sk-ant-secret", routerCfg.Providers[1].AuthHeaderValue, "anthropic value is the raw API key") + assert.Equal(t, "x-api-key", routerCfg.Providers[1].AuthHeaderName, "anthropic uses x-api-key header") + assert.Equal(t, []string{"grp-ops"}, routerCfg.Providers[1].AllowedGroupIDs, "anthropic inherits policyOps' source groups") + assert.Equal(t, []string{"claude-opus-4-7"}, routerCfg.Providers[1].Models, "anthropic's configured model ID must reach its route") + + assert.Equal(t, middlewareIDLLMLimitCheck, mws[2].ID, + "limit_check sits between router and identity_inject so deny paths skip header-stamp work") + assert.Equal(t, rpservice.MiddlewareSlotOnRequest, mws[2].Slot, "limit_check runs on_request") + + assert.Equal(t, middlewareIDLLMIdentityInject, mws[3].ID, "fourth middleware is identity inject") + assert.Equal(t, rpservice.MiddlewareSlotOnRequest, mws[3].Slot, "identity inject runs on_request") + assert.True(t, mws[3].CanMutate, "identity inject must carry CanMutate=true so its HeadersAdd / HeadersRemove pass the framework's mutation gate") + require.NotEmpty(t, mws[3].ConfigJSON, "identity inject config JSON must be populated even when no provider needs injection") + + assert.Equal(t, middlewareIDLLMGuardrail, mws[4].ID, "fifth middleware is the guardrail") + assert.Equal(t, rpservice.MiddlewareSlotOnRequest, mws[4].Slot, "guardrail runs on_request") + require.NotEmpty(t, mws[4].ConfigJSON, "guardrail config JSON must be populated") + + assert.Equal(t, middlewareIDLLMLimitRecord, mws[5].ID, + "limit_record sits FIRST in the response section so it RUNS LAST at runtime — needs cost_meter + response_parser to have stamped tokens / cost first") + assert.Equal(t, rpservice.MiddlewareSlotOnResponse, mws[5].Slot, "limit_record runs on_response") + + assert.Equal(t, middlewareIDCostMeter, mws[6].ID, "seventh middleware is the cost meter") + assert.Equal(t, rpservice.MiddlewareSlotOnResponse, mws[6].Slot, "cost meter runs on_response") + assert.Equal(t, []byte("{}"), mws[6].ConfigJSON, "cost meter carries an explicit empty config") + + assert.Equal(t, middlewareIDLLMResponseParser, mws[7].ID, "eighth middleware is the response parser") + assert.Equal(t, rpservice.MiddlewareSlotOnResponse, mws[7].Slot, "response parser runs on_response") +} + +func TestSynthesizeServices_NoSettings_ReturnsNil(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + expectSynthBaseInputs(mockStore, ctx, nil, nil, nil, nil) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + assert.Empty(t, services, "missing settings row must yield no synth") +} + +func TestSynthesizeServices_NoProviders_ReturnsNil(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), []*types.Provider{}, nil, nil) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + assert.Empty(t, services, "settings present but no providers must yield no synth") +} + +func TestSynthesizeServices_DisabledProvider_NoService(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + provider := newSynthTestProvider() + provider.Enabled = false + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{provider}, nil, nil) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + assert.Empty(t, services, "disabled provider must not synthesise a service") +} + +func TestSynthesizeServices_DisabledPolicy_NoService(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + provider := newSynthTestProvider() + policy := newSynthTestPolicy(provider.ID, "grp-eng", "") + policy.Enabled = false + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{provider}, []*types.Policy{policy}, nil) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + assert.Empty(t, services, "disabled policy must not trigger synthesis") +} + +func TestSynthesizeServices_RouterConfigOrdering(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + first := newSynthTestProvider() + first.ID = "prov-first" + first.CreatedAt = time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + + second := newSynthTestProvider() + second.ID = "prov-second" + second.ProviderID = "anthropic_api" + second.UpstreamURL = "https://api.anthropic.com" + second.APIKey = "sk-ant" + second.CreatedAt = time.Date(2026, 3, 1, 0, 0, 0, 0, time.UTC) + + third := newSynthTestProvider() + third.ID = "prov-third" + third.ProviderID = "mistral_api" + third.UpstreamURL = "https://api.mistral.ai" + third.APIKey = "sk-mistral" + third.CreatedAt = time.Date(2026, 2, 1, 0, 0, 0, 0, time.UTC) + + policy := newSynthTestPolicy(first.ID, "grp-eng", "") + policy.DestinationProviderIDs = []string{first.ID, second.ID, third.ID} + + // Pass providers in shuffled order to confirm the synth sorts them. + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{second, first, third}, + []*types.Policy{policy}, []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var routerCfg routerConfig + for _, m := range mws { + if m.ID == middlewareIDLLMRouter { + require.NoError(t, json.Unmarshal(m.ConfigJSON, &routerCfg)) + break + } + } + require.Len(t, routerCfg.Providers, 3, "all three providers must be in the router config") + assert.Equal(t, first.ID, routerCfg.Providers[0].ID, "providers ordered by created_at; first is earliest") + assert.Equal(t, third.ID, routerCfg.Providers[1].ID, "second is mid") + assert.Equal(t, second.ID, routerCfg.Providers[2].ID, "third is latest") +} + +func TestSynthesizeServices_PolicyCheckConfig_UnionsSourceGroups(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + provider := newSynthTestProvider() + + policyA := newSynthTestPolicy(provider.ID, "grp-eng", "") + policyA.ID = "pol-a" + policyA.SourceGroups = []string{"grp-eng", "grp-shared"} + policyB := newSynthTestPolicy(provider.ID, "grp-ops", "") + policyB.ID = "pol-b" + policyB.SourceGroups = []string{"grp-ops", "grp-shared"} + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{provider}, + []*types.Policy{policyA, policyB}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var routerCfg routerConfig + for _, m := range mws { + if m.ID == middlewareIDLLMRouter { + require.NoError(t, json.Unmarshal(m.ConfigJSON, &routerCfg)) + break + } + } + require.Len(t, routerCfg.Providers, 1, "single provider authorised by both policies") + assert.Equal(t, provider.ID, routerCfg.Providers[0].ID) + assert.Equal(t, []string{"grp-eng", "grp-ops", "grp-shared"}, routerCfg.Providers[0].AllowedGroupIDs, + "source groups must be unioned and sorted; the duplicate grp-shared collapses") +} + +func TestSynthesizeServices_OrphanProvider_HasEmptyAllowedGroups(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + authorised := newSynthTestProvider() + authorised.ID = "prov-authed" + + orphan := newSynthTestProvider() + orphan.ID = "prov-orphan" + orphan.ProviderID = "anthropic_api" + orphan.UpstreamURL = "https://api.anthropic.com" + orphan.APIKey = "sk-ant" + orphan.CreatedAt = time.Date(2026, 4, 1, 0, 0, 0, 0, time.UTC) + + // Policy authorises the first provider only. + policy := newSynthTestPolicy(authorised.ID, "grp-eng", "") + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{authorised, orphan}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var routerCfg routerConfig + for _, m := range mws { + if m.ID == middlewareIDLLMRouter { + require.NoError(t, json.Unmarshal(m.ConfigJSON, &routerCfg)) + break + } + } + + // Orphan providers are dropped from the router config entirely. + // The router treats an empty AllowedGroupIDs as a catch-all (right + // default for non-agent-network targets, wrong default here), so + // we don't ship them at all. Peers attempting to call models only + // the orphan claims see model_not_routable; peers calling models + // shared with the authorised provider get routed there. + require.Len(t, routerCfg.Providers, 1, "only the authorised provider reaches the router") + assert.Equal(t, authorised.ID, routerCfg.Providers[0].ID, + "authorised provider must be in router config") + assert.Equal(t, []string{"grp-eng"}, routerCfg.Providers[0].AllowedGroupIDs, + "authorised provider inherits the policy's source groups") +} + +// TestSynthesizeServices_IdentityInject_LiteLLM pins that a LiteLLM +// provider lands in the identity-inject middleware's config with the +// catalog-defined LiteLLM headers, while a non-LiteLLM provider does +// not. Together they prove the middleware is a no-op for accounts that +// don't use LiteLLM and stamps identity for those that do. +func TestSynthesizeServices_IdentityInject_LiteLLM(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + openai := newSynthTestProvider() + openai.ID = "prov-openai" + + litellm := newSynthTestProvider() + litellm.ID = "prov-litellm" + litellm.ProviderID = "litellm_proxy" + litellm.UpstreamURL = "https://litellm.acme.example.com" + litellm.APIKey = "sk-llm-master" + litellm.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC) + + policyOpenAI := newSynthTestPolicy(openai.ID, "grp-eng", "") + policyOpenAI.ID = "pol-openai" + policyLiteLLM := newSynthTestPolicy(litellm.ID, "grp-eng", "") + policyLiteLLM.ID = "pol-litellm" + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{openai, litellm}, + []*types.Policy{policyOpenAI, policyLiteLLM}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var injectCfg identityInjectConfig + for _, m := range mws { + if m.ID == middlewareIDLLMIdentityInject { + require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg)) + break + } + } + require.Len(t, injectCfg.Providers, 1, + "only providers whose catalog entry declares IdentityInjection should appear in the inject config") + entry := injectCfg.Providers[0] + assert.Equal(t, litellm.ID, entry.ProviderID, + "the LiteLLM provider must be the one identity-stamped, not the OpenAI direct provider") + require.NotNil(t, entry.HeaderPair, "LiteLLM uses the HeaderPair shape") + assert.Nil(t, entry.JSONMetadata, "shapes are mutually exclusive — JSONMetadata must be nil for HeaderPair providers") + assert.Equal(t, "x-litellm-end-user-id", entry.HeaderPair.EndUserIDHeader, + "end-user-id header must come from the catalog entry's IdentityInjection block") + assert.Equal(t, "x-litellm-tags", entry.HeaderPair.TagsHeader) +} + +// TestSynthesizeServices_IdentityInject_Bifrost_OperatorOverrides +// covers the customizable HeaderPair contract. The Bifrost catalog +// entry sets HeaderPair.Customizable=true with x-bf-dim-* defaults +// (placeholders surfaced by the dashboard, NOT authoritative at +// synth time). The wire header names that actually land on the +// inject middleware config come from the provider record's +// IdentityHeaderUserID / IdentityHeaderGroups fields verbatim. This +// lets operators pick between Bifrost's two attribution paths +// (always-on x-bf-lh-* logs metadata vs. label-declared x-bf-dim-* +// telemetry) per provider record without code changes. +// +// Three sub-cases under one fixture: full override, partial +// override (user kept, groups disabled), and ParserID empty so the +// proxy falls back to URL sniffing. +func TestSynthesizeServices_IdentityInject_Bifrost_OperatorOverrides(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + bifrost := newSynthTestProvider() + bifrost.ID = "prov-bifrost" + bifrost.ProviderID = "bifrost" + bifrost.UpstreamURL = "https://bifrost.acme.example.com/openai/v1" + bifrost.APIKey = "sk-bf-key" + bifrost.IdentityHeaderUserID = "x-bf-lh-netbird_user_id" + bifrost.IdentityHeaderGroups = "x-bf-lh-netbird_groups" + bifrost.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC) + + policy := newSynthTestPolicy(bifrost.ID, "grp-eng", "") + policy.ID = "pol-bifrost" + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{bifrost}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var injectCfg identityInjectConfig + for _, m := range mws { + if m.ID == middlewareIDLLMIdentityInject { + require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg)) + break + } + } + require.Len(t, injectCfg.Providers, 1, + "single bifrost catalog entry → one inject config target — operator's URL path picks the parser, not the catalog id") + + entry := injectCfg.Providers[0] + assert.Equal(t, bifrost.ID, entry.ProviderID) + require.NotNil(t, entry.HeaderPair, "Bifrost uses HeaderPair shape") + assert.Equal(t, "x-bf-lh-netbird_user_id", entry.HeaderPair.EndUserIDHeader, + "operator-set IdentityHeaderUserID overrides the catalog's x-bf-dim- placeholder — proves the Customizable flag actually swaps the source of truth") + assert.Equal(t, "x-bf-lh-netbird_groups", entry.HeaderPair.TagsHeader, + "operator-set IdentityHeaderGroups overrides the catalog's x-bf-dim- placeholder") + assert.False(t, entry.HeaderPair.TagsInBody, + "body-inject flags stay catalog-owned — Bifrost reads identity from headers, body inject would be a no-op") + assert.False(t, entry.HeaderPair.EndUserIDInBody) +} + +// TestSynthesizeServices_IdentityInject_Bifrost_PartialDisable proves +// that clearing one of the IdentityHeader* fields disables stamping +// for THAT dimension only, leaving the other dimension active. +// Critical because the customizable contract says "empty = disabled +// for that dimension"; if the synth path silently fell back to the +// catalog default for an empty operator value, operators couldn't +// turn off groups while keeping user id (or vice versa). +func TestSynthesizeServices_IdentityInject_Bifrost_PartialDisable(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + bifrost := newSynthTestProvider() + bifrost.ID = "prov-bifrost" + bifrost.ProviderID = "bifrost" + bifrost.UpstreamURL = "https://bifrost.acme.example.com/openai/v1" + bifrost.APIKey = "sk-bf-key" + bifrost.IdentityHeaderUserID = "x-bf-lh-netbird_user_id" + bifrost.IdentityHeaderGroups = "" // operator explicitly disabled groups + bifrost.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC) + + policy := newSynthTestPolicy(bifrost.ID, "grp-eng", "") + policy.ID = "pol-bifrost" + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{bifrost}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var injectCfg identityInjectConfig + for _, m := range mws { + if m.ID == middlewareIDLLMIdentityInject { + require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg)) + break + } + } + require.Len(t, injectCfg.Providers, 1) + entry := injectCfg.Providers[0] + require.NotNil(t, entry.HeaderPair, "user-id header is still set so the rule fires") + assert.Equal(t, "x-bf-lh-netbird_user_id", entry.HeaderPair.EndUserIDHeader) + assert.Empty(t, entry.HeaderPair.TagsHeader, + "groups header must be empty — operator cleared it; the inject middleware no-ops on empty header names so groups are NOT stamped") +} + +// TestSynthesizeServices_IdentityInject_Cloudflare_OperatorOverrides +// covers the JSONMetadata customizable contract: Cloudflare's +// catalog entry sets JSONMetadata.Customizable=true with +// netbird_user_id / netbird_groups defaults that the dashboard +// surfaces as placeholders. The actual JSON keys that land inside +// the cf-aig-metadata header come from the provider record's +// IdentityHeaderUserID / IdentityHeaderGroups fields. Reuses the +// same fields HeaderPair customizable does — the dimensions +// (user identity, groups) match; only the wire encoding (JSON key +// vs HTTP header name) differs. +func TestSynthesizeServices_IdentityInject_Cloudflare_OperatorOverrides(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + cf := newSynthTestProvider() + cf.ID = "prov-cf" + cf.ProviderID = "cloudflare_ai_gateway" + cf.UpstreamURL = "https://gateway.ai.cloudflare.com/v1/acct-xyz/my-gateway/openai" + cf.APIKey = "cf-aig-token" + cf.IdentityHeaderUserID = "team_member" + cf.IdentityHeaderGroups = "team_groups" + cf.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC) + + policy := newSynthTestPolicy(cf.ID, "grp-eng", "") + policy.ID = "pol-cf" + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{cf}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var injectCfg identityInjectConfig + for _, m := range mws { + if m.ID == middlewareIDLLMIdentityInject { + require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg)) + break + } + } + require.Len(t, injectCfg.Providers, 1) + entry := injectCfg.Providers[0] + require.NotNil(t, entry.JSONMetadata, "Cloudflare uses JSONMetadata shape — single header carrying a JSON object") + assert.Nil(t, entry.HeaderPair, "shapes are mutually exclusive") + assert.Equal(t, "cf-aig-metadata", entry.JSONMetadata.Header, + "the wire header is catalog-owned (cf-aig-metadata) — operator can rename the JSON keys but not the header itself") + assert.Equal(t, "team_member", entry.JSONMetadata.UserKey, + "operator-set IdentityHeaderUserID overrides the catalog's netbird_user_id default — proves the JSONMetadata Customizable flag swaps the source of truth like HeaderPair already does") + assert.Equal(t, "team_groups", entry.JSONMetadata.GroupsKey, + "operator-set IdentityHeaderGroups overrides the catalog's netbird_groups default") +} + +// TestSynthesizeServices_IdentityInject_Portkey_NotCustomizable +// is the JSONMetadata negative case: Portkey's catalog entry leaves +// Customizable=false because Portkey's analytics dashboard reserves +// "_user" and "groups" as fixed JSON keys. An operator-set +// IdentityHeader* on a Portkey provider record must NOT override +// those keys, or Portkey's per-user filters silently break. +func TestSynthesizeServices_IdentityInject_Portkey_NotCustomizable(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + portkey := newSynthTestProvider() + portkey.ID = "prov-portkey" + portkey.ProviderID = "portkey" + portkey.UpstreamURL = "https://api.portkey.ai/v1" + portkey.APIKey = "portkey-account-key" + // Operator set these — but portkey's catalog entry has + // JSONMetadata.Customizable=false, so synth must IGNORE them + // and stick with the catalog's _user / groups defaults. + portkey.IdentityHeaderUserID = "should-be-ignored" + portkey.IdentityHeaderGroups = "should-be-ignored-too" + portkey.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC) + + policy := newSynthTestPolicy(portkey.ID, "grp-eng", "") + policy.ID = "pol-portkey" + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{portkey}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var injectCfg identityInjectConfig + for _, m := range mws { + if m.ID == middlewareIDLLMIdentityInject { + require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg)) + break + } + } + require.Len(t, injectCfg.Providers, 1) + entry := injectCfg.Providers[0] + require.NotNil(t, entry.JSONMetadata) + assert.Equal(t, "_user", entry.JSONMetadata.UserKey, + "Portkey's reserved JSON key must hold — Customizable=false on the catalog blocks the operator's override fields") + assert.Equal(t, "groups", entry.JSONMetadata.GroupsKey, + "same fixed-schema guarantee for the groups dimension") +} + +// TestSynthesizeServices_IdentityInject_Vercel pins Vercel AI +// Gateway's wiring: HeaderPair shape with fixed wire names dictated +// by Vercel's Custom Reporting API (ai-reporting-user / +// ai-reporting-tags). Customizable=false on the catalog entry, so +// the synth path takes the catalog values verbatim and ignores any +// IdentityHeader* fields the operator might have set. Renaming +// these headers would just silently disable attribution — Vercel's +// reporting endpoint only matches the canonical names — so the +// fixed contract is the right semantic. +func TestSynthesizeServices_IdentityInject_Vercel(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + vercel := newSynthTestProvider() + vercel.ID = "prov-vercel" + vercel.ProviderID = "vercel_ai_gateway" + vercel.UpstreamURL = "https://ai-gateway.vercel.sh/v1" + vercel.APIKey = "vrc-team-key" + // Operator set these — they MUST be ignored because Vercel's + // catalog entry is non-customizable. Renaming the headers on + // the wire would defeat Vercel's reporting endpoint. + vercel.IdentityHeaderUserID = "should-be-ignored" + vercel.IdentityHeaderGroups = "should-be-ignored-too" + vercel.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC) + + policy := newSynthTestPolicy(vercel.ID, "grp-eng", "") + policy.ID = "pol-vercel" + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{vercel}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var injectCfg identityInjectConfig + for _, m := range mws { + if m.ID == middlewareIDLLMIdentityInject { + require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg)) + break + } + } + require.Len(t, injectCfg.Providers, 1) + entry := injectCfg.Providers[0] + require.NotNil(t, entry.HeaderPair, "Vercel uses HeaderPair shape — separate ai-reporting-user / ai-reporting-tags headers, not a JSON blob") + assert.Nil(t, entry.JSONMetadata, "shapes are mutually exclusive") + assert.Equal(t, "ai-reporting-user", entry.HeaderPair.EndUserIDHeader, + "end-user-id header must be Vercel's canonical ai-reporting-user — renaming would silently disable attribution at Vercel's Custom Reporting endpoint") + assert.Equal(t, "ai-reporting-tags", entry.HeaderPair.TagsHeader, + "tags header must be Vercel's canonical ai-reporting-tags for the same reason") + assert.False(t, entry.HeaderPair.TagsInBody, + "Vercel reads from headers — body inject would be a LiteLLM-specific belt-and-suspenders unneeded here") + assert.False(t, entry.HeaderPair.EndUserIDInBody) +} + +// TestSynthesizeServices_IdentityInject_OpenRouter pins OpenRouter's +// wiring: HeaderPair shape with body-only injection. OpenRouter's +// per-user attribution is the OpenAI-standard `user` body field — +// there's no header path and no groups dimension at all. The catalog +// entry sets EndUserIDInBody=true with empty header names; the inject +// middleware writes user identity into the request body but stamps +// nothing on the header surface. Customizable=false so any operator +// IdentityHeader* fields are ignored. +// +// Also asserts the static ExtraHeaders surface: operators provide +// their app URL and display name on the provider record (HTTP-Referer +// and X-OpenRouter-Title), and these land on every upstream request +// so OpenRouter's app rankings / analytics attribute correctly. +func TestSynthesizeServices_IdentityInject_OpenRouter(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + openrouter := newSynthTestProvider() + openrouter.ID = "prov-openrouter" + openrouter.ProviderID = "openrouter" + openrouter.UpstreamURL = "https://openrouter.ai/api/v1" + openrouter.APIKey = "sk-or-v1-acme" + // These would only apply if the catalog entry was Customizable; + // it isn't, so they must be IGNORED. + openrouter.IdentityHeaderUserID = "should-be-ignored" + openrouter.IdentityHeaderGroups = "should-be-ignored-too" + openrouter.ExtraValues = map[string]string{ + "HTTP-Referer": "https://acme.example/agents", + "X-OpenRouter-Title": "Acme Agents", + } + openrouter.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC) + + policy := newSynthTestPolicy(openrouter.ID, "grp-eng", "") + policy.ID = "pol-openrouter" + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{openrouter}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var injectCfg identityInjectConfig + for _, m := range mws { + if m.ID == middlewareIDLLMIdentityInject { + require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg)) + break + } + } + require.Len(t, injectCfg.Providers, 1) + entry := injectCfg.Providers[0] + require.NotNil(t, entry.HeaderPair, "OpenRouter uses HeaderPair shape — body-inject is the only branch active") + assert.Empty(t, entry.HeaderPair.EndUserIDHeader, + "OpenRouter does not document a header path for per-user identity; the inject must NOT stamp a header here. Customizable=false means operator IdentityHeader* fields are ignored.") + assert.Empty(t, entry.HeaderPair.TagsHeader, + "OpenRouter has no per-request groups / tags dimension — the tags header MUST stay empty") + assert.True(t, entry.HeaderPair.EndUserIDInBody, + "OpenRouter's only per-user attribution path is the OpenAI-standard `user` body field — body inject is the load-bearing piece for this provider") + assert.False(t, entry.HeaderPair.TagsInBody, + "no tags dimension at all → no tags-in-body either") + + // ExtraHeaders carry the operator-typed app URL + display name to + // OpenRouter's app rankings. The synth must echo BOTH static + // header values with the operator's typed strings. + require.Len(t, entry.ExtraHeaders, 2, + "both ExtraHeaders the catalog declares should land on the inject config when the operator filled in values") + byName := map[string]string{} + for _, h := range entry.ExtraHeaders { + byName[h.Name] = h.Value + } + assert.Equal(t, "https://acme.example/agents", byName["HTTP-Referer"], + "HTTP-Referer is OpenRouter's primary app identifier — must round-trip the operator-typed URL verbatim") + assert.Equal(t, "Acme Agents", byName["X-OpenRouter-Title"], + "X-OpenRouter-Title surfaces as the app's display name in OpenRouter's rankings — must round-trip operator's chosen string") +} + +// TestSynthesizeServices_IdentityInject_NonCustomizable_UsesCatalog +// is the LiteLLM-style negative case: when the catalog entry does +// NOT flag HeaderPair as Customizable, the catalog defaults are +// authoritative and any IdentityHeader* values on the provider +// record are ignored. Without this guard, an operator who set those +// fields on a non-Bifrost provider could accidentally break the +// gateway's wire protocol (LiteLLM only honours x-litellm-end-user- +// id; renaming it would silently drop spend tracking). +func TestSynthesizeServices_IdentityInject_NonCustomizable_UsesCatalog(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + litellm := newSynthTestProvider() + litellm.ID = "prov-litellm" + litellm.ProviderID = "litellm_proxy" + litellm.UpstreamURL = "https://litellm.acme.example.com" + litellm.APIKey = "sk-llm-master" + // Operator set these — but litellm_proxy's catalog entry has + // HeaderPair.Customizable=false, so the synth path must IGNORE + // these and fall back to the catalog defaults. + litellm.IdentityHeaderUserID = "x-bf-lh-should-be-ignored" + litellm.IdentityHeaderGroups = "x-bf-lh-should-be-ignored-too" + litellm.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC) + + policy := newSynthTestPolicy(litellm.ID, "grp-eng", "") + policy.ID = "pol-litellm" + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{litellm}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var injectCfg identityInjectConfig + for _, m := range mws { + if m.ID == middlewareIDLLMIdentityInject { + require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg)) + break + } + } + require.Len(t, injectCfg.Providers, 1) + entry := injectCfg.Providers[0] + require.NotNil(t, entry.HeaderPair) + assert.Equal(t, "x-litellm-end-user-id", entry.HeaderPair.EndUserIDHeader, + "Customizable=false on the catalog entry must hold — operator IdentityHeader* fields cannot rename a fixed wire protocol's headers") + assert.Equal(t, "x-litellm-tags", entry.HeaderPair.TagsHeader, + "Customizable=false on the catalog entry must hold for tags too") +} + +func TestSynthesizeServices_GuardrailMerge_AllowlistUnion_LimitsRestrictive(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + provider := newSynthTestProvider() + + guardrailA := &types.Guardrail{ + ID: "g-a", + AccountID: testAccountID, + Checks: types.GuardrailChecks{ + ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true, Models: []string{"gpt-5.4-mini"}}, + }, + } + guardrailB := &types.Guardrail{ + ID: "g-b", + AccountID: testAccountID, + Checks: types.GuardrailChecks{ + ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true, Models: []string{"gpt-5.4-pro"}}, + }, + } + + policyA := newSynthTestPolicy(provider.ID, "grp-a", guardrailA.ID) + policyA.ID = "pol-a" + policyB := newSynthTestPolicy(provider.ID, "grp-b", guardrailB.ID) + policyB.ID = "pol-b" + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{provider}, + []*types.Policy{policyA, policyB}, + []*types.Guardrail{guardrailA, guardrailB}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var guardrailJSON []byte + for _, m := range mws { + if m.ID == middlewareIDLLMGuardrail { + guardrailJSON = m.ConfigJSON + break + } + } + require.NotEmpty(t, guardrailJSON, "guardrail middleware config JSON must be present") + + var cfg guardrailConfig + require.NoError(t, json.Unmarshal(guardrailJSON, &cfg), "guardrail config must unmarshal cleanly") + assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ModelAllowlist, + "model allowlist union must keep both models") +} + +func TestSynthesizeServices_BackfillsMissingSessionKeys(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + provider := newSynthTestProvider() + provider.SessionPrivateKey = "" + provider.SessionPublicKey = "" + policy := newSynthTestPolicy(provider.ID, "grp-eng", "") + + mockStore.EXPECT(). + GetAgentNetworkSettings(ctx, store.LockingStrengthNone, testAccountID). + Return(newSynthTestSettings(), nil) + mockStore.EXPECT(). + GetAccountAgentNetworkProviders(ctx, store.LockingStrengthNone, testAccountID). + Return([]*types.Provider{provider}, nil) + mockStore.EXPECT(). + GetAccountAgentNetworkPolicies(ctx, store.LockingStrengthNone, testAccountID). + Return([]*types.Policy{policy}, nil) + // Backfill must persist the new keys before synthesising. + mockStore.EXPECT(). + SaveAgentNetworkProvider(ctx, gomock.Any()). + DoAndReturn(func(_ context.Context, p *types.Provider) error { + require.NotEmpty(t, p.SessionPrivateKey, "backfill must populate private key") + require.NotEmpty(t, p.SessionPublicKey, "backfill must populate public key") + return nil + }) + mockStore.EXPECT(). + GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, testAccountID). + Return([]*types.Guardrail{}, nil) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1, "synthesis must complete after backfill") + assert.NotEmpty(t, services[0].SessionPrivateKey, "synthesised service inherits the freshly-minted private key") + assert.NotEmpty(t, services[0].SessionPublicKey, "synthesised service inherits the freshly-minted public key") +} + +func TestSynthesizeServices_HTTPUpstream_KeepsExplicitPort(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + provider := newSynthTestProvider() + provider.UpstreamURL = "http://internal-llm.lan:8080" + policy := newSynthTestPolicy(provider.ID, "grp-eng", "") + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{provider}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var routerCfg routerConfig + for _, m := range mws { + if m.ID == middlewareIDLLMRouter { + require.NoError(t, json.Unmarshal(m.ConfigJSON, &routerCfg)) + break + } + } + require.Len(t, routerCfg.Providers, 1) + assert.Equal(t, "http", routerCfg.Providers[0].UpstreamScheme, "scheme follows the upstream URL") + assert.Equal(t, "internal-llm.lan:8080", routerCfg.Providers[0].UpstreamHost, + "explicit port travels with host so the router rewrite carries an authority the proxy can dial") +} + +func TestSynthesizeServices_UpstreamURLPath_FlowsToRouter(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + // Provider configured with a path-prefixed upstream — common for + // OpenAI-compatible endpoints behind corporate gateways. The path + // is the router's disambiguator when two providers claim the same + // model, so it must round-trip through buildRouterConfigJSON with + // the trailing slash trimmed. + provider := newSynthTestProvider() + provider.UpstreamURL = "https://corp.example.com/openai/" + policy := newSynthTestPolicy(provider.ID, "grp-eng", "") + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{provider}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + mws := services[0].Targets[0].Options.Middlewares + var routerCfg routerConfig + for _, m := range mws { + if m.ID == middlewareIDLLMRouter { + require.NoError(t, json.Unmarshal(m.ConfigJSON, &routerCfg)) + break + } + } + require.Len(t, routerCfg.Providers, 1) + assert.Equal(t, "corp.example.com", routerCfg.Providers[0].UpstreamHost, "host should drop the path") + assert.Equal(t, "/openai", routerCfg.Providers[0].UpstreamPath, + "upstream path must be carried so the router can disambiguate same-model providers; trailing slash trimmed for stable string-prefix matching") +} + +func TestSynthesizeServices_UnknownProviderID_FailsClosed(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + provider := newSynthTestProvider() + provider.ProviderID = "nonexistent_provider" + policy := newSynthTestPolicy(provider.ID, "grp-eng", "") + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{provider}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + _, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.Error(t, err, "synthesis must fail when the catalog can't resolve the provider id") + assert.Contains(t, err.Error(), "unknown catalog id", "error must surface the misconfiguration") +} + +func TestSynthesizeServices_EmptyAPIKey_FailsClosed(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + provider := newSynthTestProvider() + provider.APIKey = "" + policy := newSynthTestPolicy(provider.ID, "grp-eng", "") + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{provider}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + _, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.Error(t, err, "synthesis must refuse a provider with no api key") + assert.Contains(t, err.Error(), "no api key", "error must surface the missing credential") +} diff --git a/management/server/agentnetwork/types/accesslog.go b/management/server/agentnetwork/types/accesslog.go new file mode 100644 index 000000000..4ffb52e89 --- /dev/null +++ b/management/server/agentnetwork/types/accesslog.go @@ -0,0 +1,117 @@ +package types + +import ( + "time" + + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// AgentNetworkAccessLog is the dedicated, flattened agent-network access-log +// row. Unlike the shared reverse-proxy AccessLogEntry (which kept LLM data in +// an opaque metadata JSON blob), the LLM dimensions live in first-class, +// indexed columns so the access-log surface can filter server-side by +// user / group / provider / model / decision. +type AgentNetworkAccessLog struct { + ID string `gorm:"primaryKey"` + AccountID string `gorm:"index"` + ServiceID string `gorm:"index"` + Timestamp time.Time `gorm:"index"` + UserID string `gorm:"index"` + SourceIP string + Method string + Host string + Path string `gorm:"type:text"` + Duration time.Duration + StatusCode int `gorm:"index"` + AuthMethod string + BytesUpload int64 + BytesDownload int64 + + // Flattened LLM dimensions (queryable). Sourced from proxy metadata keys. + Provider string `gorm:"index"` // vendor, e.g. "openai" (llm.provider) + Model string `gorm:"index"` // llm.model + SessionID string `gorm:"index"` // llm.session_id — groups a conversation / coding session + ResolvedProviderID string `gorm:"index"` // llm.resolved_provider_id + SelectedPolicyID string `gorm:"index"` // llm.selected_policy_id + Decision string `gorm:"index"` // llm_policy.decision (allow/deny) + DenyReason string // llm_policy.reason (raw code, mapped in the UI) + InputTokens int64 + OutputTokens int64 + TotalTokens int64 + CostUSD float64 + Stream bool + + // Prompt capture. Only populated when prompt collection is enabled + // (account master switch AND policy guardrail). Heavy free text. + RequestPrompt string `gorm:"type:text"` + ResponseCompletion string `gorm:"type:text"` + + CreatedAt time.Time + + // GroupIDs is the authorising group ids for this entry, hydrated from the + // group child table on read. Not a column. + GroupIDs []string `gorm:"-"` +} + +// TableName keeps agent-network access logs in their own table, separate from +// the reverse-proxy AccessLogEntry table. +func (AgentNetworkAccessLog) TableName() string { return "agent_network_access_log" } + +// ToAPIResponse renders the flattened entry as the API representation. +func (a *AgentNetworkAccessLog) ToAPIResponse() api.AgentNetworkAccessLog { + out := api.AgentNetworkAccessLog{ + Id: a.ID, + ServiceId: a.ServiceID, + Timestamp: a.Timestamp, + StatusCode: a.StatusCode, + DurationMs: int(a.Duration.Milliseconds()), + InputTokens: a.InputTokens, + OutputTokens: a.OutputTokens, + TotalTokens: a.TotalTokens, + CostUsd: a.CostUSD, + Stream: &a.Stream, + } + + out.UserId = strPtr(a.UserID) + out.SourceIp = strPtr(a.SourceIP) + out.Method = strPtr(a.Method) + out.Host = strPtr(a.Host) + out.Path = strPtr(a.Path) + out.Provider = strPtr(a.Provider) + out.Model = strPtr(a.Model) + out.SessionId = strPtr(a.SessionID) + out.ResolvedProviderId = strPtr(a.ResolvedProviderID) + out.SelectedPolicyId = strPtr(a.SelectedPolicyID) + out.Decision = strPtr(a.Decision) + out.DenyReason = strPtr(a.DenyReason) + out.RequestPrompt = strPtr(a.RequestPrompt) + out.ResponseCompletion = strPtr(a.ResponseCompletion) + + if len(a.GroupIDs) > 0 { + groups := a.GroupIDs + out.GroupIds = &groups + } + return out +} + +// strPtr returns a pointer to s, or nil when s is empty — so empty optional +// fields are omitted from the JSON rather than serialised as "". +func strPtr(s string) *string { + if s == "" { + return nil + } + return &s +} + +// AgentNetworkAccessLogGroup is the normalised many-to-many row linking a log +// entry to one authorising group, so the access-log endpoint can filter by +// group with a simple `group_id IN (...)` join instead of substring-matching a +// CSV column. +type AgentNetworkAccessLogGroup struct { + LogID string `gorm:"primaryKey"` + GroupID string `gorm:"primaryKey;index"` + AccountID string `gorm:"index"` +} + +// TableName names the access-log group child table. +func (AgentNetworkAccessLogGroup) TableName() string { return "agent_network_access_log_group" } diff --git a/management/server/agentnetwork/types/accesslogfilter.go b/management/server/agentnetwork/types/accesslogfilter.go new file mode 100644 index 000000000..2906e7afa --- /dev/null +++ b/management/server/agentnetwork/types/accesslogfilter.go @@ -0,0 +1,213 @@ +package types + +import ( + "math" + "net/http" + "strconv" + "strings" + "time" + + "github.com/netbirdio/netbird/shared/management/status" +) + +const ( + // AccessLogDefaultPageSize is the default number of records per page. + AccessLogDefaultPageSize = 50 + // AccessLogMaxPageSize is the maximum number of records allowed per page. + AccessLogMaxPageSize = 100 + + accessLogDefaultSortBy = "timestamp" + accessLogDefaultSortOrder = "desc" + + // usageOverviewDefaultLookback bounds an unbounded usage-overview query so + // it never aggregates an account's entire history into memory. + usageOverviewDefaultLookback = 90 * 24 * time.Hour + // usageOverviewMaxRange caps how far back an explicit range may reach. + usageOverviewMaxRange = 366 * 24 * time.Hour +) + +// ApplyUsageOverviewBounds bounds a missing or over-wide date range so the +// in-memory usage aggregation can't load an account's full usage history. An +// absent range defaults to the last usageOverviewDefaultLookback; a range wider +// than usageOverviewMaxRange is clamped from the (possibly defaulted) end. +func (f *AgentNetworkAccessLogFilter) ApplyUsageOverviewBounds(now time.Time) { + end := now + if f.EndDate != nil { + end = *f.EndDate + } + f.EndDate = &end + if f.StartDate == nil { + start := end.Add(-usageOverviewDefaultLookback) + f.StartDate = &start + return + } + if end.Sub(*f.StartDate) > usageOverviewMaxRange { + start := end.Add(-usageOverviewMaxRange) + f.StartDate = &start + } +} + +// accessLogSortFields maps the API sort_by values to their database columns. +var accessLogSortFields = map[string]string{ + "timestamp": "timestamp", + "model": "model", + "provider": "provider", + "status_code": "status_code", + "duration": "duration", + "cost_usd": "cost_usd", + "total_tokens": "total_tokens", + "user_id": "user_id", + "decision": "decision", +} + +// AgentNetworkAccessLogFilter holds pagination, filtering and sorting +// parameters for the agent-network access-log listing. Group / provider / +// model are multi-valued (the UI uses multi-select; an entry matches when it +// matches any selected value). +type AgentNetworkAccessLogFilter struct { + Page int + PageSize int + + SortBy string + SortOrder string + + Search *string // log id, host, path, model, user email/name + UserID *string // exact user id (the dashboard sends the picked user's id) + SessionID *string // exact session id — groups one conversation / coding session + GroupIDs []string // authorising group ids (match any) + ProviderIDs []string // resolved provider ids (match any) + Models []string // models (match any) + Decision *string // policy decision (allow/deny) + PathPrefix *string // request path prefix (path LIKE 'prefix%') + StartDate *time.Time // timestamp >= start_date + EndDate *time.Time // timestamp <= end_date +} + +// ParseFromRequest fills the filter from the request query parameters. It +// returns a validation error when a supplied start_date / end_date is present +// but not valid RFC3339: silently dropping a malformed date would broaden the +// query (and, for the usage overview, fall back to the default window). +func (f *AgentNetworkAccessLogFilter) ParseFromRequest(r *http.Request) error { + q := r.URL.Query() + + f.Page = parseAccessLogPositiveInt(q.Get("page"), 1) + f.PageSize = min(parseAccessLogPositiveInt(q.Get("page_size"), AccessLogDefaultPageSize), AccessLogMaxPageSize) + + f.SortBy = parseAccessLogSortField(q.Get("sort_by")) + f.SortOrder = parseAccessLogSortOrder(q.Get("sort_order")) + + f.Search = parseAccessLogOptionalString(q.Get("search")) + f.UserID = parseAccessLogOptionalString(q.Get("user_id")) + f.SessionID = parseAccessLogOptionalString(q.Get("session_id")) + f.Decision = parseAccessLogOptionalString(q.Get("decision")) + f.PathPrefix = parseAccessLogOptionalString(q.Get("path")) + // Multi-value filters accept either repeated params (?group_id=a&group_id=b) + // or a single comma-separated value (?group_id=a,b) so both the OpenAPI + // array form and the dashboard's single-value query builder work. + f.GroupIDs = splitMultiValue(q["group_id"]) + f.ProviderIDs = splitMultiValue(q["provider_id"]) + f.Models = splitMultiValue(q["model"]) + + var err error + if f.StartDate, err = parseAccessLogOptionalRFC3339(q.Get("start_date")); err != nil { + return status.Errorf(status.InvalidArgument, "invalid start_date: %v", err) + } + if f.EndDate, err = parseAccessLogOptionalRFC3339(q.Get("end_date")); err != nil { + return status.Errorf(status.InvalidArgument, "invalid end_date: %v", err) + } + return nil +} + +// GetSortColumn returns the database column for the active sort field. +func (f *AgentNetworkAccessLogFilter) GetSortColumn() string { + if col, ok := accessLogSortFields[f.SortBy]; ok { + return col + } + return accessLogSortFields[accessLogDefaultSortBy] +} + +// GetSortOrder returns the normalised sort order ("ASC"/"DESC"). +func (f *AgentNetworkAccessLogFilter) GetSortOrder() string { + if strings.EqualFold(f.SortOrder, "asc") { + return "ASC" + } + return "DESC" +} + +// GetLimit returns the page size, defaulting/clamping when unset. +func (f *AgentNetworkAccessLogFilter) GetLimit() int { + if f.PageSize <= 0 { + return AccessLogDefaultPageSize + } + return min(f.PageSize, AccessLogMaxPageSize) +} + +// GetOffset returns the zero-based row offset for the active page. Page is +// user-controlled, so the multiplication is guarded against int overflow. +func (f *AgentNetworkAccessLogFilter) GetOffset() int { + limit := f.GetLimit() + if f.Page <= 1 || limit <= 0 { + return 0 + } + if f.Page-1 > math.MaxInt/limit { + return math.MaxInt - (math.MaxInt % limit) + } + return (f.Page - 1) * limit +} + +func parseAccessLogPositiveInt(s string, def int) int { + if v, err := strconv.Atoi(strings.TrimSpace(s)); err == nil && v > 0 { + return v + } + return def +} + +func parseAccessLogSortField(s string) string { + if _, ok := accessLogSortFields[s]; ok { + return s + } + return accessLogDefaultSortBy +} + +func parseAccessLogSortOrder(s string) string { + if strings.EqualFold(s, "asc") { + return "asc" + } + return accessLogDefaultSortOrder +} + +func parseAccessLogOptionalString(s string) *string { + if s = strings.TrimSpace(s); s != "" { + return &s + } + return nil +} + +func parseAccessLogOptionalRFC3339(s string) (*time.Time, error) { + if s = strings.TrimSpace(s); s == "" { + return nil, nil //nolint:nilnil // not provided: no value and no error + } + t, err := time.Parse(time.RFC3339, s) + if err != nil { + return nil, err + } + return &t, nil +} + +// splitMultiValue flattens repeated query params and comma-separated values +// into a single trimmed, blank-free list. Returns nil when nothing remains so +// callers can skip the filter entirely. +func splitMultiValue(values []string) []string { + out := make([]string, 0, len(values)) + for _, raw := range values { + for _, v := range strings.Split(raw, ",") { + if v = strings.TrimSpace(v); v != "" { + out = append(out, v) + } + } + } + if len(out) == 0 { + return nil + } + return out +} diff --git a/management/server/agentnetwork/types/budgetrule.go b/management/server/agentnetwork/types/budgetrule.go new file mode 100644 index 000000000..6f02a5b92 --- /dev/null +++ b/management/server/agentnetwork/types/budgetrule.go @@ -0,0 +1,106 @@ +package types + +import ( + "time" + + "github.com/rs/xid" + + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// AccountBudgetRule is an account-level, limit-only rule bound to groups +// and/or users. It mirrors the policy budget experience without any routing: +// it carries the same cap shape as a policy (PolicyLimits) but never selects a +// provider. Rules apply across policies as an always-on ceiling — every +// applicable rule binds (min-wins), so a rule can only tighten a caller's +// effective limit, never loosen it. +// +// TargetGroups matches when it intersects the caller's groups; TargetUsers +// binds a specific user directly. Empty TargetGroups and TargetUsers means the +// rule applies to every caller (the account-wide default). +type AccountBudgetRule struct { + ID string `gorm:"primaryKey"` + AccountID string `gorm:"index"` + Name string + Enabled bool + TargetGroups []string `gorm:"serializer:json;column:target_groups"` + TargetUsers []string `gorm:"serializer:json;column:target_users"` + Limits PolicyLimits `gorm:"serializer:json;column:limits"` + + CreatedAt time.Time + UpdatedAt time.Time +} + +// TableName puts budget rules in their own table. +func (AccountBudgetRule) TableName() string { return "agent_network_budget_rules" } + +// NewAccountBudgetRule returns a new rule with a freshly minted ID. +func NewAccountBudgetRule(accountID string) *AccountBudgetRule { + now := time.Now().UTC() + return &AccountBudgetRule{ + ID: "ainbud_" + xid.New().String(), + AccountID: accountID, + Enabled: true, + CreatedAt: now, + UpdatedAt: now, + } +} + +// Copy returns a deep copy of the rule, including its target slices. +func (r *AccountBudgetRule) Copy() *AccountBudgetRule { + c := *r + c.TargetGroups = append([]string(nil), r.TargetGroups...) + c.TargetUsers = append([]string(nil), r.TargetUsers...) + return &c +} + +// EventMeta renders the rule for the activity log. +func (r *AccountBudgetRule) EventMeta() map[string]any { + return map[string]any{ + "name": r.Name, + "enabled": r.Enabled, + } +} + +// FromAPIRequest applies the request payload onto the receiver. +func (r *AccountBudgetRule) FromAPIRequest(req *api.AgentNetworkBudgetRuleRequest) { + r.Name = req.Name + if req.Enabled != nil { + r.Enabled = *req.Enabled + } + if req.TargetGroups != nil { + r.TargetGroups = append([]string(nil), (*req.TargetGroups)...) + } else { + r.TargetGroups = []string{} + } + if req.TargetUsers != nil { + r.TargetUsers = append([]string(nil), (*req.TargetUsers)...) + } else { + r.TargetUsers = []string{} + } + r.Limits = limitsFromAPI(req.Limits) +} + +// ToAPIResponse renders the rule as the API representation. +func (r *AccountBudgetRule) ToAPIResponse() *api.AgentNetworkBudgetRule { + groups := r.TargetGroups + if groups == nil { + groups = []string{} + } + users := r.TargetUsers + if users == nil { + users = []string{} + } + created := r.CreatedAt + updated := r.UpdatedAt + return &api.AgentNetworkBudgetRule{ + Id: r.ID, + Name: r.Name, + Enabled: r.Enabled, + TargetGroups: groups, + TargetUsers: users, + Limits: limitsToAPI(r.Limits), + CreatedAt: &created, + UpdatedAt: &updated, + } +} diff --git a/management/server/agentnetwork/types/consumption.go b/management/server/agentnetwork/types/consumption.go new file mode 100644 index 000000000..5295b5570 --- /dev/null +++ b/management/server/agentnetwork/types/consumption.go @@ -0,0 +1,69 @@ +package types + +import "time" + +// ConsumptionDimension classifies which kind of identity a consumption +// row counts against. The proxy-side enforcement layer ticks one row +// per dimension per request — typically one user row plus one group +// row. +type ConsumptionDimension string + +const ( + // DimensionUser counts tokens / spend for a single end user. The + // dim_id column carries the netbird user id (or peer.ID when the + // caller is a tunnel-peer principal). + DimensionUser ConsumptionDimension = "user" + // DimensionGroup counts tokens / spend for a single source group + // across every member of that group. The dim_id column carries + // the netbird group id. + DimensionGroup ConsumptionDimension = "group" +) + +// Consumption is a per-dimension token + USD counter for a fixed +// aligned window. The (account, dim_kind, dim_id, window_seconds, +// window_start) tuple is the primary key; rows are rolled forward by +// the proxy's post-flight RecordLLMUsage path on every request. +// +// The same dim_id (e.g. a group id) gets one row per distinct +// window_seconds length in scope across the account's policies, +// because two policies with different window lengths read independent +// counters even though they share the dimension. Two policies with +// identical window_seconds on the same dimension share one counter +// (correct: their caps are checked against the same shared bucket). +type Consumption struct { + AccountID string `gorm:"primaryKey;type:varchar(255)"` + DimensionKind ConsumptionDimension `gorm:"primaryKey;type:varchar(16);column:dim_kind"` + DimensionID string `gorm:"primaryKey;type:varchar(255);column:dim_id"` + WindowSeconds int64 `gorm:"primaryKey;column:window_seconds"` + WindowStartUTC time.Time `gorm:"primaryKey;column:window_start_utc"` + TokensInput int64 `gorm:"column:tokens_input"` + TokensOutput int64 `gorm:"column:tokens_output"` + CostUSD float64 `gorm:"column:cost_usd"` + UpdatedAt time.Time +} + +// TableName forces a stable name independent of GORM's pluraliser. +func (Consumption) TableName() string { return "agent_network_consumption" } + +// ConsumptionKey identifies a single consumption counter within an account: +// the (dim_kind, dim_id, window_seconds, window_start) part of the row's +// primary key. Used to batch-read and batch-increment many counters for one +// request in a single store round-trip / transaction. +type ConsumptionKey struct { + Kind ConsumptionDimension + DimID string + WindowSeconds int64 + WindowStartUTC time.Time +} + +// WindowStart returns the aligned UTC start of the window of length +// windowSeconds that contains t. Aligned to the unix epoch so the +// same bucket boundary is computed deterministically across processes. +func WindowStart(t time.Time, windowSeconds int64) time.Time { + if windowSeconds <= 0 { + return t.UTC() + } + step := windowSeconds * int64(time.Second) + bucketed := t.UTC().UnixNano() / step * step + return time.Unix(0, bucketed).UTC() +} diff --git a/management/server/agentnetwork/types/consumption_test.go b/management/server/agentnetwork/types/consumption_test.go new file mode 100644 index 000000000..596cc158c --- /dev/null +++ b/management/server/agentnetwork/types/consumption_test.go @@ -0,0 +1,141 @@ +package types + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +// TestWindowStart_AlignedToUnixEpoch is the multi-node-convergence +// guarantee: any two proxies computing WindowStart(now, s) for the +// same s must land on the same boundary. The implementation aligns +// to the unix epoch (UTC) rather than local time, calendar weeks, or +// process start time — none of which are shared across nodes. +// +// Table covers the load-bearing window lengths (5m, 1h, 24h, 30d) +// plus a few odd values that still need to align cleanly. +func TestWindowStart_AlignedToUnixEpoch(t *testing.T) { + cases := []struct { + name string + instant time.Time + windowSeconds int64 + want time.Time + }{ + { + name: "5m window — drops seconds inside the bucket", + instant: time.Date(2026, 5, 6, 13, 47, 23, 0, time.UTC), + windowSeconds: 300, + want: time.Date(2026, 5, 6, 13, 45, 0, 0, time.UTC), + }, + { + name: "1h window — drops minutes / seconds, keeps the hour", + instant: time.Date(2026, 5, 6, 13, 47, 23, 0, time.UTC), + windowSeconds: 3600, + want: time.Date(2026, 5, 6, 13, 0, 0, 0, time.UTC), + }, + { + name: "24h window aligns to UTC midnight", + instant: time.Date(2026, 5, 6, 13, 47, 23, 0, time.UTC), + windowSeconds: 86_400, + want: time.Date(2026, 5, 6, 0, 0, 0, 0, time.UTC), + }, + { + name: "30d (2_592_000s) window aligns to the 30d epoch grid, not month boundaries", + instant: time.Date(2026, 5, 6, 0, 0, 0, 0, time.UTC), + windowSeconds: 2_592_000, + // 2026-05-06 UTC = 1778025600s; 1778025600 / 2592000 = 685 + // 685 * 2592000 = 1775520000s = 2026-04-07 00:00:00 UTC + want: time.Date(2026, 4, 7, 0, 0, 0, 0, time.UTC), + }, + { + name: "non-UTC input still anchors on UTC epoch boundaries", + instant: time.Date(2026, 5, 6, 13, 47, 23, 0, time.FixedZone("CEST", 2*3600)), + windowSeconds: 86_400, + // 2026-05-06 13:47:23 CEST = 11:47:23 UTC → bucket 2026-05-06 00:00:00 UTC + want: time.Date(2026, 5, 6, 0, 0, 0, 0, time.UTC), + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := WindowStart(tc.instant, tc.windowSeconds) + assert.True(t, got.Equal(tc.want), + "WindowStart(%v, %ds) = %v, want %v", tc.instant, tc.windowSeconds, got, tc.want) + }) + } +} + +// TestWindowStart_WithinWindowConverges proves the determinism +// contract: any two timestamps inside the same window land on the +// exact same boundary. Two proxy nodes serving requests 7s apart +// must agree on which counter row to upsert. +func TestWindowStart_WithinWindowConverges(t *testing.T) { + t1 := time.Date(2026, 5, 6, 14, 0, 0, 0, time.UTC) + t2 := t1.Add(7 * time.Second) + t3 := t1.Add(59*time.Minute + 59*time.Second) + + a := WindowStart(t1, 3600) + b := WindowStart(t2, 3600) + c := WindowStart(t3, 3600) + + assert.True(t, a.Equal(b), "two timestamps 7s apart in the same 1h window must align to the same boundary") + assert.True(t, a.Equal(c), "the very last second of a 1h window still lands on the SAME bucket as the first second") +} + +// TestWindowStart_AcrossWindowsDiverges is the symmetric guarantee: +// two timestamps separated by a window's worth of time MUST land on +// different boundaries. Without this, a 24h window's "rollover" +// would never reset the counter. +func TestWindowStart_AcrossWindowsDiverges(t *testing.T) { + t1 := time.Date(2026, 5, 6, 23, 59, 59, 0, time.UTC) + t2 := t1.Add(2 * time.Second) // 2026-05-07 00:00:01 + + a := WindowStart(t1, 86_400) + b := WindowStart(t2, 86_400) + assert.False(t, a.Equal(b), + "timestamps straddling a 24h-window boundary must land on different buckets — otherwise daily caps never reset") +} + +// TestWindowStart_DifferentWindowsHaveDifferentBuckets locks the +// design fork "two policies with different window_seconds on the same +// group produce independent counters". A 24h boundary at noon is NOT +// the same as the 30d boundary that contains it. +func TestWindowStart_DifferentWindowsHaveDifferentBuckets(t *testing.T) { + now := time.Date(2026, 5, 6, 12, 0, 0, 0, time.UTC) + short := WindowStart(now, 86_400) + long := WindowStart(now, 2_592_000) + assert.False(t, short.Equal(long), + "the 24h bucket and 30d bucket containing the same instant must differ — independent counters require independent keys") +} + +// TestWindowStart_SubMinuteAndMinuteAlignment locks sub-hour windows. +// A 5-minute window must align to multiples of 300s from the unix +// epoch — minute marks 0/5/10/.../55 within an hour, deterministic +// across nodes regardless of clock drift. +func TestWindowStart_SubMinuteAndMinuteAlignment(t *testing.T) { + t1 := time.Date(2026, 5, 6, 14, 12, 30, 0, time.UTC) + t2 := time.Date(2026, 5, 6, 14, 14, 59, 0, time.UTC) + t3 := time.Date(2026, 5, 6, 14, 15, 0, 0, time.UTC) + + a := WindowStart(t1, 300) + b := WindowStart(t2, 300) + c := WindowStart(t3, 300) + + assert.True(t, a.Equal(b), + "14:12:30 and 14:14:59 fall in the same 5m bucket starting at 14:10:00") + assert.True(t, a.Equal(time.Date(2026, 5, 6, 14, 10, 0, 0, time.UTC)), + "5m bucket containing 14:12 starts at 14:10 — aligned to multiples of 300s from unix epoch") + assert.False(t, a.Equal(c), + "14:15:00 is the start of the next 5m bucket — must not fold into the previous one") +} + +// TestWindowStart_ZeroWindowReturnsInputUTC covers the defensive +// path: caller hands a zero / negative window (shouldn't happen, but +// might mid-refactor). The function returns the input as UTC rather +// than dividing by zero. +func TestWindowStart_ZeroWindowReturnsInputUTC(t *testing.T) { + now := time.Date(2026, 5, 6, 12, 30, 45, 0, time.FixedZone("CEST", 2*3600)) + got := WindowStart(now, 0) + assert.True(t, got.Equal(now.UTC()), "zero window must not panic — return input as UTC") +} diff --git a/management/server/agentnetwork/types/guardrail.go b/management/server/agentnetwork/types/guardrail.go new file mode 100644 index 000000000..12edd815c --- /dev/null +++ b/management/server/agentnetwork/types/guardrail.go @@ -0,0 +1,120 @@ +package types + +import ( + "time" + + "github.com/rs/xid" + + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// GuardrailChecks is the configurable parameter set persisted with each +// guardrail. Stored as a JSON blob to keep the table flat. +type GuardrailChecks struct { + ModelAllowlist GuardrailModelAllowlist `json:"model_allowlist"` + PromptCapture GuardrailPromptCapture `json:"prompt_capture"` +} + +type GuardrailModelAllowlist struct { + Enabled bool `json:"enabled"` + Models []string `json:"models"` +} + +type GuardrailPromptCapture struct { + Enabled bool `json:"enabled"` + RedactPii bool `json:"redact_pii"` +} + +// Guardrail is an Agent Network reusable guardrail set persisted per account. +type Guardrail struct { + ID string `gorm:"primaryKey"` + AccountID string `gorm:"index"` + Name string + Description string + Checks GuardrailChecks `gorm:"serializer:json"` + CreatedAt time.Time + UpdatedAt time.Time +} + +// TableName uses an explicit name so guardrail rows live in their own +// table. +func (Guardrail) TableName() string { return "agent_network_guardrails" } + +// NewGuardrail returns a new Guardrail with a freshly minted ID. +func NewGuardrail(accountID string) *Guardrail { + now := time.Now().UTC() + return &Guardrail{ + ID: "ainguard_" + xid.New().String(), + AccountID: accountID, + Checks: GuardrailChecks{ModelAllowlist: GuardrailModelAllowlist{Models: []string{}}}, + CreatedAt: now, + UpdatedAt: now, + } +} + +// FromAPIRequest applies the request payload onto the receiver. +func (g *Guardrail) FromAPIRequest(req *api.AgentNetworkGuardrailRequest) { + g.Name = req.Name + if req.Description != nil { + g.Description = *req.Description + } + g.Checks = checksFromAPI(req.Checks) +} + +// ToAPIResponse renders the guardrail as the API representation. +func (g *Guardrail) ToAPIResponse() *api.AgentNetworkGuardrail { + created := g.CreatedAt + updated := g.UpdatedAt + return &api.AgentNetworkGuardrail{ + Id: g.ID, + Name: g.Name, + Description: g.Description, + Checks: checksToAPI(g.Checks), + CreatedAt: &created, + UpdatedAt: &updated, + } +} + +// Copy returns a deep copy of the guardrail. +func (g *Guardrail) Copy() *Guardrail { + clone := *g + if g.Checks.ModelAllowlist.Models != nil { + clone.Checks.ModelAllowlist.Models = append([]string(nil), g.Checks.ModelAllowlist.Models...) + } + return &clone +} + +// EventMeta is the audit-log payload for activity events. +func (g *Guardrail) EventMeta() map[string]any { + return map[string]any{"name": g.Name} +} + +func checksFromAPI(c api.AgentNetworkGuardrailChecks) GuardrailChecks { + models := append([]string(nil), c.ModelAllowlist.Models...) + if models == nil { + models = []string{} + } + return GuardrailChecks{ + ModelAllowlist: GuardrailModelAllowlist{ + Enabled: c.ModelAllowlist.Enabled, + Models: models, + }, + PromptCapture: GuardrailPromptCapture{ + Enabled: c.PromptCapture.Enabled, + RedactPii: c.PromptCapture.RedactPii, + }, + } +} + +func checksToAPI(c GuardrailChecks) api.AgentNetworkGuardrailChecks { + models := c.ModelAllowlist.Models + if models == nil { + models = []string{} + } + out := api.AgentNetworkGuardrailChecks{} + out.ModelAllowlist.Enabled = c.ModelAllowlist.Enabled + out.ModelAllowlist.Models = models + out.PromptCapture.Enabled = c.PromptCapture.Enabled + out.PromptCapture.RedactPii = c.PromptCapture.RedactPii + return out +} diff --git a/management/server/agentnetwork/types/policy.go b/management/server/agentnetwork/types/policy.go new file mode 100644 index 000000000..709b8149d --- /dev/null +++ b/management/server/agentnetwork/types/policy.go @@ -0,0 +1,192 @@ +package types + +import ( + "time" + + "github.com/rs/xid" + + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// Policy is an Agent Network policy persisted per account. A policy +// authorises members of SourceGroups to reach the listed +// DestinationProviderIDs under the attached GuardrailIDs and Limits. +// +// Token and budget limits live on the Policy itself (Limits field); +// guardrails carry only model allowlist and prompt capture. +type Policy struct { + ID string `gorm:"primaryKey"` + AccountID string `gorm:"index"` + Name string + Description string + Enabled bool + SourceGroups []string `gorm:"serializer:json;column:source_groups"` + DestinationProviderIDs []string `gorm:"serializer:json;column:destination_provider_ids"` + GuardrailIDs []string `gorm:"serializer:json;column:guardrail_ids"` + Limits PolicyLimits `gorm:"serializer:json;column:limits"` + + CreatedAt time.Time + UpdatedAt time.Time +} + +// PolicyLimits aggregates the token and budget caps attached directly +// to a policy. Both halves are always present; their Enabled flags +// control whether the proxy enforces them. +type PolicyLimits struct { + TokenLimit PolicyTokenLimit `json:"token_limit"` + BudgetLimit PolicyBudgetLimit `json:"budget_limit"` +} + +// PolicyTokenLimit is a token-count cap evaluated over an aligned +// window of WindowSeconds seconds. GroupCap is applied to each +// source group independently — every group in the policy's +// SourceGroups gets its own bucket of GroupCap tokens. UserCap +// applies independently to each individual user. A zero cap means +// uncapped. WindowSeconds must be at least 60 (one minute) when the +// limit is enabled. +type PolicyTokenLimit struct { + Enabled bool `json:"enabled"` + GroupCap int64 `json:"group_cap"` + UserCap int64 `json:"user_cap"` + WindowSeconds int64 `json:"window_seconds"` +} + +// PolicyBudgetLimit is a USD spend cap evaluated over an aligned +// window of WindowSeconds seconds. GroupCapUsd is applied to each +// source group independently — every group in the policy's +// SourceGroups gets its own bucket of GroupCapUsd USD. UserCapUsd +// applies independently to each individual user. A zero cap means +// uncapped. WindowSeconds must be at least 60 (one minute) when the +// limit is enabled. +type PolicyBudgetLimit struct { + Enabled bool `json:"enabled"` + GroupCapUsd float64 `json:"group_cap_usd"` + UserCapUsd float64 `json:"user_cap_usd"` + WindowSeconds int64 `json:"window_seconds"` +} + +// TableName forces a unique GORM table to avoid collision with the access +// control Policy type, which also resolves to "policies" by default. +func (Policy) TableName() string { return "agent_network_policies" } + +// NewPolicy returns a new Policy with a freshly minted ID. +func NewPolicy(accountID string) *Policy { + now := time.Now().UTC() + return &Policy{ + ID: "ainpol_" + xid.New().String(), + AccountID: accountID, + Enabled: true, + CreatedAt: now, + UpdatedAt: now, + } +} + +// FromAPIRequest applies the request payload onto the receiver. +func (p *Policy) FromAPIRequest(req *api.AgentNetworkPolicyRequest) { + p.Name = req.Name + if req.Description != nil { + p.Description = *req.Description + } + if req.Enabled != nil { + p.Enabled = *req.Enabled + } + p.SourceGroups = append([]string(nil), req.SourceGroups...) + p.DestinationProviderIDs = append([]string(nil), req.DestinationProviderIds...) + if req.GuardrailIds != nil { + p.GuardrailIDs = append([]string(nil), (*req.GuardrailIds)...) + } else { + p.GuardrailIDs = []string{} + } + if req.Limits != nil { + p.Limits = limitsFromAPI(*req.Limits) + } else { + p.Limits = PolicyLimits{} + } +} + +// ToAPIResponse renders the policy as the API representation. +func (p *Policy) ToAPIResponse() *api.AgentNetworkPolicy { + src := p.SourceGroups + if src == nil { + src = []string{} + } + dst := p.DestinationProviderIDs + if dst == nil { + dst = []string{} + } + guardrails := p.GuardrailIDs + if guardrails == nil { + guardrails = []string{} + } + created := p.CreatedAt + updated := p.UpdatedAt + return &api.AgentNetworkPolicy{ + Id: p.ID, + Name: p.Name, + Description: p.Description, + Enabled: p.Enabled, + SourceGroups: src, + DestinationProviderIds: dst, + GuardrailIds: guardrails, + Limits: limitsToAPI(p.Limits), + CreatedAt: &created, + UpdatedAt: &updated, + } +} + +// Copy returns a deep copy of the policy. +func (p *Policy) Copy() *Policy { + clone := *p + if p.SourceGroups != nil { + clone.SourceGroups = append([]string(nil), p.SourceGroups...) + } + if p.DestinationProviderIDs != nil { + clone.DestinationProviderIDs = append([]string(nil), p.DestinationProviderIDs...) + } + if p.GuardrailIDs != nil { + clone.GuardrailIDs = append([]string(nil), p.GuardrailIDs...) + } + return &clone +} + +// EventMeta is the audit-log payload for activity events. +func (p *Policy) EventMeta() map[string]any { + return map[string]any{ + "name": p.Name, + "enabled": p.Enabled, + } +} + +func limitsFromAPI(in api.AgentNetworkPolicyLimits) PolicyLimits { + return PolicyLimits{ + TokenLimit: PolicyTokenLimit{ + Enabled: in.TokenLimit.Enabled, + GroupCap: in.TokenLimit.GroupCap, + UserCap: in.TokenLimit.UserCap, + WindowSeconds: in.TokenLimit.WindowSeconds, + }, + BudgetLimit: PolicyBudgetLimit{ + Enabled: in.BudgetLimit.Enabled, + GroupCapUsd: in.BudgetLimit.GroupCapUsd, + UserCapUsd: in.BudgetLimit.UserCapUsd, + WindowSeconds: in.BudgetLimit.WindowSeconds, + }, + } +} + +func limitsToAPI(in PolicyLimits) api.AgentNetworkPolicyLimits { + return api.AgentNetworkPolicyLimits{ + TokenLimit: api.AgentNetworkPolicyTokenLimit{ + Enabled: in.TokenLimit.Enabled, + GroupCap: in.TokenLimit.GroupCap, + UserCap: in.TokenLimit.UserCap, + WindowSeconds: in.TokenLimit.WindowSeconds, + }, + BudgetLimit: api.AgentNetworkPolicyBudgetLimit{ + Enabled: in.BudgetLimit.Enabled, + GroupCapUsd: in.BudgetLimit.GroupCapUsd, + UserCapUsd: in.BudgetLimit.UserCapUsd, + WindowSeconds: in.BudgetLimit.WindowSeconds, + }, + } +} diff --git a/management/server/agentnetwork/types/provider.go b/management/server/agentnetwork/types/provider.go new file mode 100644 index 000000000..28c8a94e2 --- /dev/null +++ b/management/server/agentnetwork/types/provider.go @@ -0,0 +1,252 @@ +package types + +import ( + "fmt" + "strings" + "time" + + "github.com/rs/xid" + + "github.com/netbirdio/netbird/shared/management/http/api" + "github.com/netbirdio/netbird/util/crypt" +) + +// ProviderModel is one row in the provider's models list. The operator +// pins the per-1k input/output price for cost tracking; ID is the +// model identifier the upstream provider expects on the wire. +type ProviderModel struct { + ID string `json:"id"` + InputPer1k float64 `json:"input_per_1k"` + OutputPer1k float64 `json:"output_per_1k"` +} + +// Provider is an Agent Network AI provider record persisted per account. +// The proxy cluster fronting the account lives on the per-account +// agent-network Settings row, not on the Provider — every provider in +// an account routes through the same cluster. +type Provider struct { + ID string `gorm:"primaryKey"` + AccountID string `gorm:"index"` + ProviderID string `gorm:"index:idx_agent_network_provider"` + Name string + // UpstreamURL is the full upstream URL (e.g. https://api.openai.com) + // the operator selected. + UpstreamURL string `gorm:"column:upstream_url"` + APIKey string `gorm:"column:api_key"` + // ExtraValues holds operator-typed values for catalog-declared + // ExtraHeaders (see catalog.Provider.ExtraHeaders). Keyed by + // header name (e.g. "x-portkey-config"); a non-empty value is + // stamped on every upstream request to this provider via the + // proxy's identity-inject middleware (anti-spoof Remove + Add). + // Empty / missing keys = no header stamped. Stored as a JSON + // blob so the schema doesn't grow per-catalog-entry. + ExtraValues map[string]string `gorm:"serializer:json;column:extra_values"` + // Models is the operator's curated list of models exposed by this + // provider together with their per-1k input/output prices (USD). + // Empty means all catalog models are allowed at catalog prices. + Models []ProviderModel `gorm:"serializer:json"` + Enabled bool + // SessionPrivateKey + SessionPublicKey are the ed25519 keypair the + // synthesised reverse-proxy service uses to sign / verify session + // JWTs after a successful OIDC handshake. Generated once on + // provider create and never rotated by the manager so existing + // session cookies survive provider edits. SessionPrivateKey is + // encrypted at rest via EncryptSensitiveData / + // DecryptSensitiveData; SessionPublicKey is plain. + SessionPrivateKey string `gorm:"column:session_private_key"` + SessionPublicKey string `gorm:"column:session_public_key"` + // IdentityHeaderUserID + IdentityHeaderGroups are the operator- + // chosen wire header names for HeaderPair-style identity + // injection on catalog entries that flag the shape as + // Customizable (e.g. Bifrost, where the operator picks between + // the always-on x-bf-lh- log-metadata family and the + // label-declared x-bf-dim- telemetry family). Empty value + // disables stamping for that dimension; the inject middleware + // already no-ops on empty header names. Catalog entries with + // Customizable=false ignore these fields and use the static + // header names defined in their HeaderPairInjection block. + IdentityHeaderUserID string `gorm:"column:identity_header_user_id"` + IdentityHeaderGroups string `gorm:"column:identity_header_groups"` + CreatedAt time.Time + UpdatedAt time.Time +} + +// TableName uses an explicit name so the Agent Network provider rows live +// in their own table, separate from any future "providers"-named entity. +func (Provider) TableName() string { return "agent_network_providers" } + +// NewProvider returns a new Provider with a freshly minted ID. +func NewProvider(accountID string) *Provider { + now := time.Now().UTC() + return &Provider{ + ID: xid.New().String(), + AccountID: accountID, + CreatedAt: now, + UpdatedAt: now, + } +} + +// FromAPIRequest applies the request payload onto the receiver. The api_key +// is only overwritten when the caller provided one — empty/nil leaves the +// existing key intact, so updates can omit it. +func (p *Provider) FromAPIRequest(req *api.AgentNetworkProviderRequest) { + p.ProviderID = req.ProviderId + p.Name = req.Name + p.UpstreamURL = req.UpstreamUrl + if req.ApiKey != nil && strings.TrimSpace(*req.ApiKey) != "" { + p.APIKey = *req.ApiKey + } + if req.ExtraValues != nil { + // Replace the whole map (rather than merge) so unsetting a + // value on the dashboard actually clears it. Empty strings + // are dropped so we don't waste a row on no-op values. + next := make(map[string]string, len(*req.ExtraValues)) + for k, v := range *req.ExtraValues { + v = strings.TrimSpace(v) + if v != "" { + next[k] = v + } + } + if len(next) == 0 { + p.ExtraValues = nil + } else { + p.ExtraValues = next + } + } + p.Models = p.Models[:0] + if req.Models != nil { + for _, m := range *req.Models { + p.Models = append(p.Models, ProviderModel{ + ID: m.Id, + InputPer1k: m.InputPer1k, + OutputPer1k: m.OutputPer1k, + }) + } + } + if p.Models == nil { + p.Models = []ProviderModel{} + } + if req.Enabled != nil { + p.Enabled = *req.Enabled + } + // Identity-header overrides for catalogs flagged Customizable. + // nil pointer = "field omitted on the wire" → leave the stored + // value untouched (per the openapi description). Empty string is + // an explicit clear that disables stamping for this dimension. + if req.IdentityHeaderUserId != nil { + p.IdentityHeaderUserID = strings.TrimSpace(*req.IdentityHeaderUserId) + } + if req.IdentityHeaderGroups != nil { + p.IdentityHeaderGroups = strings.TrimSpace(*req.IdentityHeaderGroups) + } +} + +// ToAPIResponse renders the provider as the API representation. The API +// key is intentionally never surfaced. +func (p *Provider) ToAPIResponse() *api.AgentNetworkProvider { + models := make([]api.AgentNetworkProviderModel, 0, len(p.Models)) + for _, m := range p.Models { + models = append(models, api.AgentNetworkProviderModel{ + Id: m.ID, + InputPer1k: m.InputPer1k, + OutputPer1k: m.OutputPer1k, + }) + } + created := p.CreatedAt + updated := p.UpdatedAt + resp := &api.AgentNetworkProvider{ + Id: p.ID, + ProviderId: p.ProviderID, + Name: p.Name, + UpstreamUrl: p.UpstreamURL, + Models: models, + Enabled: p.Enabled, + CreatedAt: &created, + UpdatedAt: &updated, + } + if len(p.ExtraValues) > 0 { + out := make(map[string]string, len(p.ExtraValues)) + for k, v := range p.ExtraValues { + out[k] = v + } + resp.ExtraValues = &out + } + if p.IdentityHeaderUserID != "" { + v := p.IdentityHeaderUserID + resp.IdentityHeaderUserId = &v + } + if p.IdentityHeaderGroups != "" { + v := p.IdentityHeaderGroups + resp.IdentityHeaderGroups = &v + } + return resp +} + +// Copy returns a deep copy of the provider. +func (p *Provider) Copy() *Provider { + clone := *p + if p.Models != nil { + clone.Models = append([]ProviderModel(nil), p.Models...) + } + if p.ExtraValues != nil { + clone.ExtraValues = make(map[string]string, len(p.ExtraValues)) + for k, v := range p.ExtraValues { + clone.ExtraValues[k] = v + } + } + return &clone +} + +// EventMeta is the audit-log payload for activity events. +func (p *Provider) EventMeta() map[string]any { + return map[string]any{ + "name": p.Name, + "provider_id": p.ProviderID, + } +} + +// EncryptSensitiveData encrypts the upstream API key and the session +// signing key in place. +func (p *Provider) EncryptSensitiveData(enc *crypt.FieldEncrypt) error { + if enc == nil { + return nil + } + if p.APIKey != "" { + encrypted, err := enc.Encrypt(p.APIKey) + if err != nil { + return fmt.Errorf("encrypt agent network provider api key: %w", err) + } + p.APIKey = encrypted + } + if p.SessionPrivateKey != "" { + encrypted, err := enc.Encrypt(p.SessionPrivateKey) + if err != nil { + return fmt.Errorf("encrypt agent network provider session key: %w", err) + } + p.SessionPrivateKey = encrypted + } + return nil +} + +// DecryptSensitiveData decrypts the upstream API key and the session +// signing key in place. +func (p *Provider) DecryptSensitiveData(enc *crypt.FieldEncrypt) error { + if enc == nil { + return nil + } + if p.APIKey != "" { + decrypted, err := enc.Decrypt(p.APIKey) + if err != nil { + return fmt.Errorf("decrypt agent network provider api key: %w", err) + } + p.APIKey = decrypted + } + if p.SessionPrivateKey != "" { + decrypted, err := enc.Decrypt(p.SessionPrivateKey) + if err != nil { + return fmt.Errorf("decrypt agent network provider session key: %w", err) + } + p.SessionPrivateKey = decrypted + } + return nil +} diff --git a/management/server/agentnetwork/types/settings.go b/management/server/agentnetwork/types/settings.go new file mode 100644 index 000000000..d61d9deff --- /dev/null +++ b/management/server/agentnetwork/types/settings.go @@ -0,0 +1,78 @@ +package types + +import ( + "time" + + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// DefaultAccessLogRetentionDays is the retention applied to new accounts' +// agent-network access logs. Usage records are not subject to this — they are +// the long-term aggregate and are retained independently. +const DefaultAccessLogRetentionDays = 30 + +// Settings is the per-account agent-network configuration row. One +// row per account. Cluster + Subdomain are immutable once written and +// produce the public endpoint agents call (`.`). +type Settings struct { + AccountID string `gorm:"primaryKey"` + Cluster string + Subdomain string `gorm:"index:idx_agent_network_settings_cluster_subdomain"` + + // Account-level collection controls sourced by the synthesizer. + // EnableLogCollection gates the per-request access-log trail and defaults + // ON for new accounts. EnablePromptCollection is the master gate for + // request/response prompt capture (AND-gated with the policy-level + // guardrail). RedactPii enables PII redaction on captured prompts; + // effective redaction is account OR policy. + EnableLogCollection bool + EnablePromptCollection bool + RedactPii bool + + // AccessLogRetentionDays bounds how long full access-log rows are kept; a + // periodic sweep deletes older rows. <= 0 means keep indefinitely. Usage + // records are unaffected. + AccessLogRetentionDays int + + CreatedAt time.Time + UpdatedAt time.Time +} + +// TableName puts the rows in their own table to keep the agent-network +// schema cohesive. +func (Settings) TableName() string { return "agent_network_settings" } + +// Endpoint returns the bare hostname agents reach this account at: +// `.`. +func (s *Settings) Endpoint() string { + return s.Subdomain + "." + s.Cluster +} + +// ToAPIResponse renders the settings as the API representation. +func (s *Settings) ToAPIResponse() *api.AgentNetworkSettings { + created := s.CreatedAt + updated := s.UpdatedAt + retention := s.AccessLogRetentionDays + return &api.AgentNetworkSettings{ + Cluster: s.Cluster, + Subdomain: s.Subdomain, + Endpoint: s.Endpoint(), + EnableLogCollection: s.EnableLogCollection, + EnablePromptCollection: s.EnablePromptCollection, + RedactPii: s.RedactPii, + AccessLogRetentionDays: &retention, + CreatedAt: &created, + UpdatedAt: &updated, + } +} + +// FromAPIRequest applies the mutable settings fields from the request. Cluster +// and Subdomain are immutable and intentionally not touched here. +func (s *Settings) FromAPIRequest(req *api.AgentNetworkSettingsRequest) { + s.EnableLogCollection = req.EnableLogCollection + s.EnablePromptCollection = req.EnablePromptCollection + s.RedactPii = req.RedactPii + if req.AccessLogRetentionDays != nil { + s.AccessLogRetentionDays = *req.AccessLogRetentionDays + } +} diff --git a/management/server/agentnetwork/types/usage.go b/management/server/agentnetwork/types/usage.go new file mode 100644 index 000000000..dd01d4300 --- /dev/null +++ b/management/server/agentnetwork/types/usage.go @@ -0,0 +1,47 @@ +package types + +import ( + "time" +) + +// AgentNetworkUsage is the stripped, always-collected per-request usage record +// powering the Usage overview. Unlike AgentNetworkAccessLog it carries no +// request detail (host/path/source IP/prompt) — only the dimensions needed to +// aggregate and filter spend by user / group / provider / model over time. +// +// It is written unconditionally on every served agent-network request, +// independent of the account's EnableLogCollection toggle: when log collection +// is off the proxy ships a stripped, usage-only entry and management still +// records the usage row (but skips the full AgentNetworkAccessLog row). +type AgentNetworkUsage struct { + ID string `gorm:"primaryKey"` + AccountID string `gorm:"index"` + Timestamp time.Time `gorm:"index"` + UserID string `gorm:"index"` + ResolvedProviderID string `gorm:"index"` + Provider string // vendor, e.g. "openai" + Model string `gorm:"index"` + SessionID string `gorm:"index"` // llm.session_id — groups a conversation / coding session + InputTokens int64 + OutputTokens int64 + TotalTokens int64 + CostUSD float64 + CreatedAt time.Time +} + +// TableName keeps usage records in their own stripped table. Named +// distinctly (…_request_usage) to avoid colliding with any pre-existing +// agent_network_usage table in a shared database. +func (AgentNetworkUsage) TableName() string { return "agent_network_request_usage" } + +// AgentNetworkUsageGroup is the normalised many-to-many row linking a usage +// record to one authorising group, mirroring AgentNetworkAccessLogGroup so the +// usage overview can filter by group with a `group_id IN (...)` join. +type AgentNetworkUsageGroup struct { + UsageID string `gorm:"primaryKey"` + GroupID string `gorm:"primaryKey;index"` + AccountID string `gorm:"index"` +} + +// TableName names the usage group child table. +func (AgentNetworkUsageGroup) TableName() string { return "agent_network_request_usage_group" } diff --git a/management/server/agentnetwork/types/usageoverview.go b/management/server/agentnetwork/types/usageoverview.go new file mode 100644 index 000000000..658832bec --- /dev/null +++ b/management/server/agentnetwork/types/usageoverview.go @@ -0,0 +1,96 @@ +package types + +import ( + "sort" + "time" + + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// UsageGranularity is the time-bucket width for the usage overview. New values +// can be added here and handled in bucketStart without touching the store. +type UsageGranularity string + +const ( + UsageGranularityDay UsageGranularity = "day" + UsageGranularityWeek UsageGranularity = "week" + UsageGranularityMonth UsageGranularity = "month" +) + +// ParseUsageGranularity maps the API query value to a granularity, defaulting +// to day for empty/unknown input. +func ParseUsageGranularity(s string) UsageGranularity { + switch UsageGranularity(s) { + case UsageGranularityWeek: + return UsageGranularityWeek + case UsageGranularityMonth: + return UsageGranularityMonth + default: + return UsageGranularityDay + } +} + +// AgentNetworkUsageBucket is one aggregated usage time bucket. PeriodStart is +// the UTC start of the bucket as YYYY-MM-DD. +type AgentNetworkUsageBucket struct { + PeriodStart string + InputTokens int64 + OutputTokens int64 + TotalTokens int64 + CostUSD float64 +} + +// ToAPIResponse renders the bucket as the API representation. +func (b *AgentNetworkUsageBucket) ToAPIResponse() api.AgentNetworkUsageBucket { + return api.AgentNetworkUsageBucket{ + PeriodStart: b.PeriodStart, + InputTokens: b.InputTokens, + OutputTokens: b.OutputTokens, + TotalTokens: b.TotalTokens, + CostUsd: b.CostUSD, + } +} + +// bucketStart truncates t (in UTC) to the start of its bucket for the given +// granularity. Week buckets start on Monday (ISO week). +func bucketStart(t time.Time, g UsageGranularity) time.Time { + t = t.UTC() + switch g { + case UsageGranularityWeek: + // Monday-start week. time.Weekday: Sunday=0..Saturday=6. + offset := (int(t.Weekday()) + 6) % 7 + day := time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, time.UTC) + return day.AddDate(0, 0, -offset) + case UsageGranularityMonth: + return time.Date(t.Year(), t.Month(), 1, 0, 0, 0, 0, time.UTC) + default: // day + return time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, time.UTC) + } +} + +// AggregateUsageByGranularity buckets the usage rows by the requested +// granularity and returns the buckets ordered oldest-first. Aggregation is done +// in Go (rather than per-engine SQL date_trunc) so granularities stay portable +// across SQLite/Postgres/MySQL and easy to extend. +func AggregateUsageByGranularity(rows []*AgentNetworkUsage, g UsageGranularity) []*AgentNetworkUsageBucket { + byPeriod := make(map[string]*AgentNetworkUsageBucket) + for _, r := range rows { + key := bucketStart(r.Timestamp, g).Format("2006-01-02") + b := byPeriod[key] + if b == nil { + b = &AgentNetworkUsageBucket{PeriodStart: key} + byPeriod[key] = b + } + b.InputTokens += r.InputTokens + b.OutputTokens += r.OutputTokens + b.TotalTokens += r.TotalTokens + b.CostUSD += r.CostUSD + } + + out := make([]*AgentNetworkUsageBucket, 0, len(byPeriod)) + for _, b := range byPeriod { + out = append(out, b) + } + sort.Slice(out, func(i, j int) bool { return out[i].PeriodStart < out[j].PeriodStart }) + return out +} diff --git a/management/server/agentnetwork/wire_shape_test.go b/management/server/agentnetwork/wire_shape_test.go new file mode 100644 index 000000000..5aa580549 --- /dev/null +++ b/management/server/agentnetwork/wire_shape_test.go @@ -0,0 +1,109 @@ +package agentnetwork + +import ( + "context" + "encoding/json" + "testing" + + "github.com/golang/mock/gomock" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy" + rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/store" + "github.com/netbirdio/netbird/shared/management/proto" +) + +// TestSynthesizedService_WireShape locks down the proto shape that +// flows from the synthesizer through ToProtoMapping to the proxy. +// Drift between this test and what the proxy expects manifests as +// "service not matching" — the proxy receives a mapping but can't +// register an SNI/HTTP route from it. +func TestSynthesizedService_WireShape(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + provider := newSynthTestProvider() + policy := newSynthTestPolicy(provider.ID, "grp-eng", "") + + expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(), + []*types.Provider{provider}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + svc := services[0] + mapping := svc.ToProtoMapping(rpservice.Create, "test-token", proxy.OIDCValidationConfig{}) + + // Identifiers — account-scoped service ID, settings-derived domain. + assert.Equal(t, "agent-net-svc-acct-1", mapping.GetId(), "stable account-scoped virtual service ID") + assert.Equal(t, testAccountID, mapping.GetAccountId(), "account id round-trips") + assert.Equal(t, testEndpoint, mapping.GetDomain(), "domain matches settings.Endpoint() output") + + // Mode + listen port — addMapping at proxy/server.go switches on Mode. + assert.Equal(t, "http", mapping.GetMode(), "synthesised services are HTTP mode") + assert.Equal(t, int32(0), mapping.GetListenPort(), "no custom listen port for HTTP services") + + // Auth token + private/tunnel shape: agent-network endpoints authenticate + // inbound agents via ValidateTunnelPeer against AccessGroups, not OIDC. + assert.Equal(t, "test-token", mapping.GetAuthToken(), "auth token round-trips for proxy CreateProxyPeer") + assert.True(t, mapping.GetPrivate(), "synthesised services are private (tunnel-peer auth via AccessGroups)") + require.NotNil(t, mapping.GetAuth(), "auth payload carries the session key") + assert.False(t, mapping.GetAuth().GetOidc(), "OIDC is off for tunnel-auth agent-network services") + + // Path mappings — proxy/server.go::setupHTTPMapping early-returns when + // len(mapping.GetPath()) == 0, so this is a critical assertion. + require.Len(t, mapping.GetPath(), 1, "exactly one path mapping for the cluster target") + pm := mapping.GetPath()[0] + assert.Equal(t, "/", pm.GetPath(), "default path is '/'") + assert.Equal(t, "https://noop.invalid/", pm.GetTarget(), + "target URL is the placeholder; the router middleware rewrites it per request") + require.NotNil(t, pm.GetOptions(), "target options must be populated so direct_upstream + middleware chain reach the proxy") + assert.True(t, pm.GetOptions().GetDirectUpstream(), "synth targets imply direct_upstream so the proxy dials via the host stack") + assert.True(t, pm.GetOptions().GetAgentNetwork(), "agent_network flag must travel on the wire so the proxy can tag access logs") + + mws := pm.GetOptions().GetMiddlewares() + require.Len(t, mws, 8, "eight middlewares reach the proxy: request_parser, router, limit_check, identity_inject, guardrail, limit_record, cost_meter, response_parser") + + assert.Equal(t, middlewareIDLLMRequestParser, mws[0].GetId(), "first middleware id") + assert.Equal(t, proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST, mws[0].GetSlot(), "request parser slot") + + assert.Equal(t, middlewareIDLLMRouter, mws[1].GetId(), "second middleware id") + assert.Equal(t, proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST, mws[1].GetSlot(), "router slot") + require.NotEmpty(t, mws[1].GetConfigJson(), "router config must travel on the wire") + var routerCfg routerConfig + require.NoError(t, json.Unmarshal(mws[1].GetConfigJson(), &routerCfg), "router config decodes") + require.Len(t, routerCfg.Providers, 1, "the only enabled provider reaches the router") + assert.Equal(t, provider.ID, routerCfg.Providers[0].ID, "router provider id matches synth provider") + assert.Equal(t, "Bearer sk-test-key", routerCfg.Providers[0].AuthHeaderValue, + "openai catalog template substitutes the API key on the wire") + + assert.Equal(t, middlewareIDLLMLimitCheck, mws[2].GetId(), + "limit_check runs after the router so the resolved provider id is available, before identity_inject so a deny doesn't pay the header-stamp cost") + assert.Equal(t, proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST, mws[2].GetSlot()) + + assert.Equal(t, middlewareIDLLMIdentityInject, mws[3].GetId(), "fourth middleware id") + assert.Equal(t, proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST, mws[3].GetSlot(), "identity inject slot") + require.NotEmpty(t, mws[3].GetConfigJson(), "identity inject config JSON must travel on the wire") + + assert.Equal(t, middlewareIDLLMGuardrail, mws[4].GetId(), "fifth middleware id") + assert.Equal(t, proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST, mws[4].GetSlot(), "guardrail slot") + require.NotEmpty(t, mws[4].GetConfigJson(), "guardrail middleware config JSON must travel on the wire") + + assert.Equal(t, middlewareIDLLMLimitRecord, mws[5].GetId(), + "limit_record sits FIRST in the response section so it RUNS LAST at runtime — slot order on the response leg is reverse-of-slice") + assert.Equal(t, proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_RESPONSE, mws[5].GetSlot()) + + assert.Equal(t, middlewareIDCostMeter, mws[6].GetId(), "seventh middleware id") + assert.Equal(t, proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_RESPONSE, mws[6].GetSlot(), "cost meter slot") + + assert.Equal(t, middlewareIDLLMResponseParser, mws[7].GetId(), "eighth middleware id") + assert.Equal(t, proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_RESPONSE, mws[7].GetSlot(), "response parser slot") +} diff --git a/management/server/agentnetwork_budgetrule_realstack_test.go b/management/server/agentnetwork_budgetrule_realstack_test.go new file mode 100644 index 000000000..31f62b1b5 --- /dev/null +++ b/management/server/agentnetwork_budgetrule_realstack_test.go @@ -0,0 +1,126 @@ +package server + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/management/server/agentnetwork" + agenttypes "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/permissions" + "github.com/netbirdio/netbird/management/server/store" +) + +// TestAgentNetwork_BudgetRuleCRUD_RealManager is the GC-1 no-mock guard for the +// account budget-rule manager surface: real DefaultAccountManager, real store, +// real permissions. It exercises create/get/list/update/delete through the +// permission-gated manager (not the store directly) and asserts the reused +// PolicyLimits cap shape and targets survive each step. +func TestAgentNetwork_BudgetRuleCRUD_RealManager(t *testing.T) { + am, _, err := createManager(t) + require.NoError(t, err, "createManager must succeed") + ctx := context.Background() + + const ( + accountID = "agent-net-budget-acct" + adminUserID = "agent-net-budget-admin" + ) + account := newAccountWithId(ctx, accountID, adminUserID, "agent-net.test", "", "", false) + require.NoError(t, am.Store.SaveAccount(ctx, account), "SaveAccount must succeed") + + mgr := agentnetwork.NewManager(am.Store, permissions.NewManager(am.Store), am, nil) + + created, err := mgr.CreateBudgetRule(ctx, adminUserID, &agenttypes.AccountBudgetRule{ + AccountID: accountID, + Name: "eng-monthly", + Enabled: true, + TargetGroups: []string{"grp-eng"}, + TargetUsers: []string{"user-alice"}, + Limits: agenttypes.PolicyLimits{ + TokenLimit: agenttypes.PolicyTokenLimit{Enabled: true, GroupCap: 100_000, UserCap: 10_000, WindowSeconds: 2_592_000}, + BudgetLimit: agenttypes.PolicyBudgetLimit{Enabled: true, GroupCapUsd: 500, WindowSeconds: 2_592_000}, + }, + }) + require.NoError(t, err, "CreateBudgetRule must succeed") + require.NotEmpty(t, created.ID, "create must mint an ID") + + got, err := mgr.GetBudgetRule(ctx, accountID, adminUserID, created.ID) + require.NoError(t, err, "GetBudgetRule must succeed") + assert.Equal(t, "eng-monthly", got.Name, "name round-trips through the manager") + assert.Equal(t, []string{"grp-eng"}, got.TargetGroups, "target groups round-trip") + assert.Equal(t, int64(100_000), got.Limits.TokenLimit.GroupCap, "token group cap round-trips") + + list, err := mgr.GetAllBudgetRules(ctx, accountID, adminUserID) + require.NoError(t, err, "GetAllBudgetRules must succeed") + require.Len(t, list, 1, "exactly the one created rule must be listed") + + created.Limits.TokenLimit.GroupCap = 200_000 + updated, err := mgr.UpdateBudgetRule(ctx, adminUserID, created) + require.NoError(t, err, "UpdateBudgetRule must succeed") + assert.Equal(t, int64(200_000), updated.Limits.TokenLimit.GroupCap, "updated cap must persist") + + require.NoError(t, mgr.DeleteBudgetRule(ctx, accountID, adminUserID, created.ID), "DeleteBudgetRule must succeed") + _, err = mgr.GetBudgetRule(ctx, accountID, adminUserID, created.ID) + assert.Error(t, err, "get after delete must fail") +} + +// TestAgentNetwork_UpdateSettings_PreservesImmutableAndTogglesCollection is the +// GC-1 guard for UpdateSettings: it must apply the collection toggles while +// preserving the immutable Cluster/Subdomain pinned at bootstrap. +func TestAgentNetwork_UpdateSettings_PreservesImmutableAndTogglesCollection(t *testing.T) { + am, _, err := createManager(t) + require.NoError(t, err, "createManager must succeed") + ctx := context.Background() + + const ( + accountID = "agent-net-settings-acct" + adminUserID = "agent-net-settings-admin" + clusterAddr = "eu.proxy.netbird.io" + ) + account := newAccountWithId(ctx, accountID, adminUserID, "agent-net.test", "", "", false) + require.NoError(t, am.Store.SaveAccount(ctx, account), "SaveAccount must succeed") + + mgr := agentnetwork.NewManager(am.Store, permissions.NewManager(am.Store), am, nil) + + // Creating a provider bootstraps the settings row (cluster + subdomain). + _, err = mgr.CreateProvider(ctx, adminUserID, &agenttypes.Provider{ + AccountID: accountID, + ProviderID: "openai_api", + Name: "openai", + UpstreamURL: "https://api.openai.com", + APIKey: "sk-test", + Enabled: true, + Models: []agenttypes.ProviderModel{{ID: "gpt-5.4"}}, + }, clusterAddr) + require.NoError(t, err, "CreateProvider must bootstrap settings") + + before, err := mgr.GetSettings(ctx, accountID, adminUserID) + require.NoError(t, err, "GetSettings must succeed after bootstrap") + require.Equal(t, clusterAddr, before.Cluster, "cluster pinned at bootstrap") + require.NotEmpty(t, before.Subdomain, "subdomain pinned at bootstrap") + assert.False(t, before.EnablePromptCollection, "prompt collection defaults off") + + // Attempt to flip toggles AND smuggle a different cluster/subdomain — the + // immutable fields must be ignored. + updated, err := mgr.UpdateSettings(ctx, adminUserID, &agenttypes.Settings{ + AccountID: accountID, + Cluster: "attacker.cluster", + Subdomain: "evil", + EnableLogCollection: true, + EnablePromptCollection: true, + RedactPii: true, + }) + require.NoError(t, err, "UpdateSettings must succeed") + assert.Equal(t, before.Cluster, updated.Cluster, "cluster is immutable and must be preserved") + assert.Equal(t, before.Subdomain, updated.Subdomain, "subdomain is immutable and must be preserved") + assert.True(t, updated.EnableLogCollection, "log collection toggle must apply") + assert.True(t, updated.EnablePromptCollection, "prompt collection toggle must apply") + assert.True(t, updated.RedactPii, "redact toggle must apply") + + reloaded, err := am.Store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID) + require.NoError(t, err) + assert.Equal(t, before.Cluster, reloaded.Cluster, "persisted cluster unchanged") + assert.True(t, reloaded.EnablePromptCollection, "persisted prompt collection toggled on") +} diff --git a/management/server/agentnetwork_proxypeer_restart_test.go b/management/server/agentnetwork_proxypeer_restart_test.go new file mode 100644 index 000000000..0957cc555 --- /dev/null +++ b/management/server/agentnetwork_proxypeer_restart_test.go @@ -0,0 +1,199 @@ +package server + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/management/internals/modules/peers" + "github.com/netbirdio/netbird/management/server/agentnetwork" + agenttypes "github.com/netbirdio/netbird/management/server/agentnetwork/types" + nbpeer "github.com/netbirdio/netbird/management/server/peer" + "github.com/netbirdio/netbird/management/server/permissions" + "github.com/netbirdio/netbird/management/server/store" + "github.com/netbirdio/netbird/management/server/types" +) + +// TestAgentNetwork_ProxyRestart_PropagatesNewPeerAndDropsStale is the no-mock +// regression guard for the bug the user reported: restarting the proxy creates +// a fresh embedded peer with a NEW WireGuard public key (the proxy generates +// the keypair on every startup at proxy/internal/roundtrip/netbird.go:312). +// The PRIOR embedded peer record is never deleted on management, so the +// account accumulates a stale peer holding a stale CGNAT IP. Other peers +// in the account either keep routing to the dead IP, or — if synth DNS +// picks the wrong record — never see the new IP at all. +// +// What this test exercises (no mocks): +// - real SQLite test store +// - real DefaultAccountManager, network-map controller, peer-update channels +// - real peers.Manager.CreateProxyPeer path (the very method the proxy +// invokes over gRPC on every startup) +// - real agentnetwork.Manager + synth chain so the client receives a +// concrete DNS record that must point at the LATEST proxy peer. +// +// Pre-fix expected behavior (red): two embedded peers exist after the +// "restart"; the synth DNS record points at the stale one; the client +// receives an update reflecting the new peer but the old one lingers. +// Post-fix expected behavior (green): exactly one embedded peer exists +// after restart (with the new key) AND the client's network map carries +// the synth DNS pointing at that new peer's CGNAT IP. +func TestAgentNetwork_ProxyRestart_PropagatesNewPeerAndDropsStale(t *testing.T) { + am, updateManager, err := createManager(t) + require.NoError(t, err, "createManager must succeed") + ctx := context.Background() + + const ( + accountID = "an-restart-acct" + adminUserID = "an-restart-admin" + groupAID = "an-restart-grp-A" + clusterAddr = "eu.proxy.netbird.io" + clientKey = "BhRPtynAAYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8=" + // Two different proxy pubkeys — the "before" and "after" of a + // proxy-process restart with fresh-keypair generation. + proxyKey1 = "Aaaaa1aaaaYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8=" + proxyKey2 = "Bbbbb2bbbbYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8=" + ) + + // --- Account scaffold --- + account := newAccountWithId(ctx, accountID, adminUserID, "an-restart.test", "", "", false) + require.NoError(t, am.Store.SaveAccount(ctx, account)) + + clientPeer := &nbpeer.Peer{ + Key: clientKey, + Name: "an-restart-client", + DNSLabel: "an-restart-client", + Meta: nbpeer.PeerSystemMeta{Hostname: "an-restart-client", GoOS: "linux", WtVersion: "development"}, + } + addedClient, _, _, _, err := am.AddPeer(ctx, "", "", adminUserID, clientPeer, false) + require.NoError(t, err, "AddPeer for client must succeed") + require.NoError(t, am.MarkPeerConnected(ctx, clientKey, accountID, time.Now().UnixNano(), &types.NetworkMap{}), + "MarkPeerConnected for the client peer must succeed (affected-peer fan-out skips disconnected peers)") + + // Place the client in group A so the synth policy reaches it. + account, err = am.Store.GetAccount(ctx, accountID) + require.NoError(t, err) + account.Groups[groupAID] = &types.Group{ID: groupAID, Name: "groupA", Peers: []string{addedClient.ID}} + require.NoError(t, am.Store.SaveAccount(ctx, account), "SaveAccount must persist group A") + + // --- Real peers + agent-network managers --- + permMgr := permissions.NewManager(am.Store) + peersMgr := peers.NewManager(am.Store, permMgr) + peersMgr.SetAccountManager(am) + peersMgr.SetNetworkMapController(am.networkMapController) + agentMgr := agentnetwork.NewManager(am.Store, permMgr, am, nil) + + // Subscribe BEFORE any state-mutating call so we don't lose the update + // that contains the synth DNS record. + clientCh := updateManager.CreateChannel(ctx, addedClient.ID) + t.Cleanup(func() { updateManager.CloseChannel(ctx, addedClient.ID) }) + drain(clientCh) + + // --- First proxy startup: register peer key K1, then mark it + // connected. In production the proxy follows CreateProxyPeer with the + // regular sync stream which lands on MarkPeerConnected; the synth DNS + // path filters out peers that aren't Connected (types/account.go:323), + // so without this step no DNS record would be emitted. + require.NoError(t, peersMgr.CreateProxyPeer(ctx, accountID, proxyKey1, clusterAddr), + "first CreateProxyPeer (proxy startup) must succeed") + + peer1ID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey1) + require.NoError(t, err, "proxy peer for K1 must be persisted after CreateProxyPeer") + require.NotEmpty(t, peer1ID) + + require.NoError(t, am.MarkPeerConnected(ctx, proxyKey1, accountID, time.Now().UnixNano(), &types.NetworkMap{}), + "MarkPeerConnected for K1 must succeed") + + account, err = am.Store.GetAccount(ctx, accountID) + require.NoError(t, err) + proxyIP1 := account.Peers[peer1ID].IP.String() + require.NotEmpty(t, proxyIP1, "K1 must have an assigned overlay IP") + + // --- Provider + policy. CreateProvider / CreatePolicy trigger the + // agentnetwork reconcile which runs UpdateAccountPeers; the resulting + // NetworkMap delivered to the client carries the synth DNS record + // pointing at K1's IP. --- + provider, err := agentMgr.CreateProvider(ctx, adminUserID, &agenttypes.Provider{ + AccountID: accountID, + ProviderID: "openai_api", + Name: "openai-test", + UpstreamURL: "https://api.openai.com", + APIKey: "sk-test-key", + Enabled: true, + Models: []agenttypes.ProviderModel{{ID: "gpt-5.4"}}, + }, clusterAddr) + require.NoError(t, err, "CreateProvider must succeed") + + _, err = agentMgr.CreatePolicy(ctx, adminUserID, &agenttypes.Policy{ + AccountID: accountID, + Name: "p1", + Enabled: true, + SourceGroups: []string{groupAID}, + DestinationProviderIDs: []string{provider.ID}, + }) + require.NoError(t, err, "CreatePolicy must succeed") + + settings, err := am.Store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID) + require.NoError(t, err) + fqdn := settings.Endpoint() + + rdata1 := awaitZoneRData(clientCh, clusterAddr, fqdn, true) + require.Equal(t, proxyIP1, rdata1, + "client must receive a synth DNS record pointing at K1's overlay IP after the synth path runs") + drain(clientCh) + + // --- Proxy restart: NEW keypair K2, same account, same cluster --- + require.NoError(t, peersMgr.CreateProxyPeer(ctx, accountID, proxyKey2, clusterAddr), + "second CreateProxyPeer (proxy restart with fresh keypair) must succeed") + + peer2ID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey2) + require.NoError(t, err, "proxy peer for K2 must be persisted after restart") + require.NotEmpty(t, peer2ID) + + require.NoError(t, am.MarkPeerConnected(ctx, proxyKey2, accountID, time.Now().UnixNano(), &types.NetworkMap{}), + "MarkPeerConnected for K2 must succeed") + + // In production the agent's sync stream pulls a fresh NetworkMap as + // part of its normal reconcile cadence; in this isolated test + // MarkPeerConnected's affected-peer fan-out can race the channel-side + // buffer in a way that swallows the synth-DNS-bearing update before + // our await reads it. Trigger an explicit account-wide fan-out so the + // assertion below tests what production actually delivers, not the + // in-test buffer race. + am.UpdateAccountPeers(ctx, accountID, types.UpdateReason{Resource: types.UpdateResourcePeer, Operation: types.UpdateOperationUpdate}) + + account, err = am.Store.GetAccount(ctx, accountID) + require.NoError(t, err) + proxyIP2 := account.Peers[peer2ID].IP.String() + require.NotEmpty(t, proxyIP2, "K2 must have an assigned overlay IP") + require.NotEqual(t, proxyIP1, proxyIP2, "K2 must get a different overlay IP than K1 (sanity)") + + // CRITICAL ASSERTION 1: K1 must no longer be in the store. The SqlStore + // returns ("", nil) for a missing key rather than NotFound, so assert + // on the returned ID being empty. + staleID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey1) + require.NoError(t, err, "GetPeerIDByKey for a missing peer must not error") + assert.Empty(t, staleID, + "stale embedded proxy peer K1 must be removed when a new embedded peer registers for the same (account, cluster); pre-fix this assertion fails because management never cleans up the prior peer record") + + // CRITICAL ASSERTION 2: exactly one embedded proxy peer remains, and it + // is K2. + account, err = am.Store.GetAccount(ctx, accountID) + require.NoError(t, err) + embeddedKeys := []string{} + for _, p := range account.Peers { + if p.ProxyMeta.Embedded { + embeddedKeys = append(embeddedKeys, p.Key) + } + } + assert.Equal(t, []string{proxyKey2}, embeddedKeys, + "after a proxy restart exactly one embedded proxy peer should remain — the one with the new key K2") + + // CRITICAL ASSERTION 3: the synth DNS record the client receives now + // points at K2's IP, not K1's. + rdata2 := awaitZoneRData(clientCh, clusterAddr, fqdn, true) + assert.Equal(t, proxyIP2, rdata2, + "after proxy restart, the client's synth DNS record must point at the NEW embedded peer's IP, not the stale K1 IP") +} diff --git a/management/server/agentnetwork_realstack_test.go b/management/server/agentnetwork_realstack_test.go new file mode 100644 index 000000000..563abd047 --- /dev/null +++ b/management/server/agentnetwork_realstack_test.go @@ -0,0 +1,212 @@ +package server + +import ( + "context" + "net/netip" + "testing" + "time" + + "github.com/miekg/dns" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + networkmap "github.com/netbirdio/netbird/management/internals/controllers/network_map" + "github.com/netbirdio/netbird/management/server/agentnetwork" + agenttypes "github.com/netbirdio/netbird/management/server/agentnetwork/types" + nbpeer "github.com/netbirdio/netbird/management/server/peer" + "github.com/netbirdio/netbird/management/server/permissions" + "github.com/netbirdio/netbird/management/server/store" + "github.com/netbirdio/netbird/management/server/types" + nbproto "github.com/netbirdio/netbird/shared/management/proto" +) + +// TestAgentNetwork_ProviderCRUD_FansOutToProxyAndClientPeers is the no-mock +// integration test for the live propagation path: a provider/policy mutation +// through the real agentnetwork.Manager triggers the real +// DefaultAccountManager.UpdateAccountPeers, which runs the real network-map +// controller (including AN-2b's injectAllProxyPolicies), and a network map is +// computed and fanned out to BOTH the embedded proxy peer and the client peer. +// +// Unlike the synthesizer/reconcile unit tests, nothing here is mocked: real +// SQLite store, real account manager + network-map controller, real +// agentnetwork manager, real peer update channels. The client peer's delivered +// map is asserted to actually carry the synth DNS surface, and provider +// create/delete are exercised end to end. +func TestAgentNetwork_ProviderCRUD_FansOutToProxyAndClientPeers(t *testing.T) { + am, updateManager, err := createManager(t) + require.NoError(t, err, "createManager must succeed") + ctx := context.Background() + + const ( + accountID = "agent-net-acct-1" + adminUserID = "agent-net-admin-1" + groupAID = "agent-net-grp-A" + clusterAddr = "eu.proxy.netbird.io" + clientKey = "BhRPtynAAYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8=" + proxyPeerID = "agent-net-proxy-peer-1" + proxyPeerKey = "/yF0+vCfv+mRR5k0dca0TrGdO/oiNeAI58gToZm5NyI=" + proxyIP = "100.64.0.99" + ) + + account := newAccountWithId(ctx, accountID, adminUserID, "agent-net.test", "", "", false) + require.NoError(t, am.Store.SaveAccount(ctx, account), "SaveAccount must succeed") + + // Real client peer through the production AddPeer path. + clientPeer := &nbpeer.Peer{ + Key: clientKey, + Name: "agent-net-client", + DNSLabel: "agent-net-client", + Meta: nbpeer.PeerSystemMeta{Hostname: "agent-net-client", GoOS: "linux", WtVersion: "development"}, + } + addedClient, _, _, _, err := am.AddPeer(ctx, "", "", adminUserID, clientPeer, false) + require.NoError(t, err, "AddPeer must add the client peer") + + // Inject a connected embedded proxy peer + put the client in the source group. + account, err = am.Store.GetAccount(ctx, accountID) + require.NoError(t, err) + account.Peers[proxyPeerID] = &nbpeer.Peer{ + ID: proxyPeerID, + AccountID: accountID, + Key: proxyPeerKey, + IP: netip.MustParseAddr(proxyIP), + Status: &nbpeer.PeerStatus{Connected: true}, + ProxyMeta: nbpeer.ProxyMeta{Embedded: true, Cluster: clusterAddr}, + DNSLabel: "agent-net-proxy", + } + account.Groups[groupAID] = &types.Group{ID: groupAID, Name: "groupA", Peers: []string{addedClient.ID}} + require.NoError(t, am.Store.SaveAccount(ctx, account), "SaveAccount must persist proxy peer + group") + + // Subscribe to BOTH peers' update channels — this is how we observe the + // real fan-out. + clientCh := updateManager.CreateChannel(ctx, addedClient.ID) + proxyCh := updateManager.CreateChannel(ctx, proxyPeerID) + t.Cleanup(func() { + updateManager.CloseChannel(ctx, addedClient.ID) + updateManager.CloseChannel(ctx, proxyPeerID) + }) + drain(clientCh) + drain(proxyCh) + + // Real agentnetwork manager wired to the real account manager. proxyController + // is nil (no gRPC cluster fan-out here) — the reconcile still fires + // UpdateAccountPeers, which is the path under test. + agentMgr := agentnetwork.NewManager(am.Store, permissions.NewManager(am.Store), am, nil) + + provider, err := agentMgr.CreateProvider(ctx, adminUserID, &agenttypes.Provider{ + AccountID: accountID, + ProviderID: "openai_api", + Name: "openai-test", + UpstreamURL: "https://api.openai.com", + APIKey: "sk-test-key", + Enabled: true, + Models: []agenttypes.ProviderModel{{ID: "gpt-5.4"}}, + }, clusterAddr) + require.NoError(t, err, "CreateProvider must succeed") + + policy, err := agentMgr.CreatePolicy(ctx, adminUserID, &agenttypes.Policy{ + AccountID: accountID, + Name: "p1", + Enabled: true, + SourceGroups: []string{groupAID}, + DestinationProviderIDs: []string{provider.ID}, + }) + require.NoError(t, err, "CreatePolicy must succeed") + + settings, err := am.Store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID) + require.NoError(t, err) + fqdn := settings.Endpoint() + + // Both peers must receive a fan-out. The provider-create reconcile fires + // before the policy exists (synth service then has no AccessGroups, so no + // zone), and the async update buffer can collapse/reorder updates — so we + // poll until the client's delivered map actually carries the synth record. + rdata := awaitZoneRData(clientCh, clusterAddr, fqdn, true) + assert.Equal(t, proxyIP, rdata, + "client peer's delivered network map must contain the synth DNS record pointing at the embedded proxy peer") + require.True(t, awaitUpdate(proxyCh), "embedded proxy peer must also receive a netmap update after create") + + // UPDATE the provider — a new model on the existing service must still + // reconcile and keep the private surface routable (the live MODIFIED path). + provider.Models = append(provider.Models, agenttypes.ProviderModel{ID: "gpt-5.4-mini"}) + _, err = agentMgr.UpdateProvider(ctx, adminUserID, provider) + require.NoError(t, err, "UpdateProvider must succeed") + assert.Equal(t, proxyIP, awaitZoneRData(clientCh, clusterAddr, fqdn, true), + "client peer must still resolve the synth record after the provider is updated") + require.True(t, awaitUpdate(proxyCh), "embedded proxy peer must also receive a netmap update after update") + + // DELETE: detach the policy first (provider is in use), then drop the + // provider. Both peers update again and the synth surface disappears. + require.NoError(t, agentMgr.DeletePolicy(ctx, accountID, adminUserID, policy.ID), "DeletePolicy must succeed") + require.NoError(t, agentMgr.DeleteProvider(ctx, accountID, adminUserID, provider.ID), "DeleteProvider must succeed") + + require.True(t, awaitUpdate(proxyCh), "embedded proxy peer must also receive a netmap update after delete") + assert.Empty(t, awaitZoneRData(clientCh, clusterAddr, fqdn, false), + "synth DNS record must be gone from the client's map after the provider is deleted") +} + +// awaitZoneRData drains the channel for up to 8s. When wantPresent is true it +// returns as soon as the synth record appears (its RData). When false it drains +// to quiescence and returns the RData of the last delivered map (expected empty +// once the provider is gone), tolerating stale buffered updates that still +// carry the zone. +func awaitZoneRData(ch <-chan *networkmap.UpdateMessage, clusterAddr, fqdn string, wantPresent bool) string { + deadline := time.After(8 * time.Second) + last := "" + for { + select { + case m := <-ch: + if m == nil { + continue + } + last = synthZoneRData(m.Update, clusterAddr, fqdn) + if wantPresent && last != "" { + return last + } + case <-time.After(750 * time.Millisecond): + return last + case <-deadline: + return last + } + } +} + +// awaitUpdate reports whether at least one update arrives within the window. +func awaitUpdate(ch <-chan *networkmap.UpdateMessage) bool { + select { + case m := <-ch: + return m != nil + case <-time.After(5 * time.Second): + return false + } +} + +// drain empties any buffered updates (e.g. from AddPeer/SaveAccount) so the +// next observation reflects the operation under test. +func drain(ch <-chan *networkmap.UpdateMessage) { + for { + select { + case <-ch: + case <-time.After(200 * time.Millisecond): + return + } + } +} + +// synthZoneRData returns the RData of the synth A record (record name == fqdn) +// inside the cluster's custom zone, or "" when absent. +func synthZoneRData(sync *nbproto.SyncResponse, clusterAddr, fqdn string) string { + if sync == nil { + return "" + } + for _, zone := range sync.GetNetworkMap().GetDNSConfig().GetCustomZones() { + if zone.GetDomain() != dns.Fqdn(clusterAddr) { + continue + } + for _, rec := range zone.GetRecords() { + if rec.GetName() == dns.Fqdn(fqdn) { + return rec.GetRData() + } + } + } + return "" +} diff --git a/management/server/http/handler.go b/management/server/http/handler.go index 0abdb854d..6e61eb129 100644 --- a/management/server/http/handler.go +++ b/management/server/http/handler.go @@ -28,6 +28,8 @@ import ( "github.com/netbirdio/netbird/management/internals/modules/zones/records" recordsManager "github.com/netbirdio/netbird/management/internals/modules/zones/records/manager" "github.com/netbirdio/netbird/management/server/account" + "github.com/netbirdio/netbird/management/server/agentnetwork" + agentnetworkhandlers "github.com/netbirdio/netbird/management/server/http/handlers/agentnetwork" "github.com/netbirdio/netbird/management/server/settings" "github.com/netbirdio/netbird/management/server/permissions" @@ -59,7 +61,7 @@ import ( ) // NewAPIHandler creates the Management service HTTP API handler registering all the available endpoints. -func NewAPIHandler(ctx context.Context, router *mux.Router, accountManager account.Manager, networksManager nbnetworks.Manager, resourceManager resources.Manager, routerManager routers.Manager, groupsManager nbgroups.Manager, LocationManager geolocation.Geolocation, authManager auth.Manager, appMetrics telemetry.AppMetrics, permissionsManager permissions.Manager, settingsManager settings.Manager, zManager zones.Manager, rManager records.Manager, networkMapController network_map.Controller, idpManager idpmanager.Manager, serviceManager service.Manager, reverseProxyDomainManager *manager.Manager, reverseProxyAccessLogsManager accesslogs.Manager, proxyGRPCServer *nbgrpc.ProxyServiceServer, trustedHTTPProxies []netip.Prefix, rateLimiter *middleware.APIRateLimiter, isValidChildAccount middleware.IsValidChildAccountFunc) (http.Handler, error) { +func NewAPIHandler(ctx context.Context, router *mux.Router, accountManager account.Manager, networksManager nbnetworks.Manager, resourceManager resources.Manager, routerManager routers.Manager, groupsManager nbgroups.Manager, LocationManager geolocation.Geolocation, authManager auth.Manager, appMetrics telemetry.AppMetrics, permissionsManager permissions.Manager, settingsManager settings.Manager, zManager zones.Manager, rManager records.Manager, networkMapController network_map.Controller, idpManager idpmanager.Manager, serviceManager service.Manager, reverseProxyDomainManager *manager.Manager, reverseProxyAccessLogsManager accesslogs.Manager, proxyGRPCServer *nbgrpc.ProxyServiceServer, trustedHTTPProxies []netip.Prefix, rateLimiter *middleware.APIRateLimiter, isValidChildAccount middleware.IsValidChildAccountFunc, agentNetworkManager agentnetwork.Manager) (http.Handler, error) { // Register bypass paths for unauthenticated endpoints if err := bypass.AddBypassPath("/api/instance"); err != nil { @@ -124,6 +126,9 @@ func NewAPIHandler(ctx context.Context, router *mux.Router, accountManager accou zonesManager.RegisterEndpoints(router, zManager) recordsManager.RegisterEndpoints(router, rManager) idp.AddEndpoints(accountManager, router) + if agentNetworkManager != nil { + agentnetworkhandlers.AddEndpoints(agentNetworkManager, router) + } instance.AddEndpoints(instanceManager, accountManager, router) instance.AddVersionEndpoint(instanceManager, router) if serviceManager != nil && reverseProxyDomainManager != nil { diff --git a/management/server/http/handlers/agentnetwork/access_log_handler.go b/management/server/http/handlers/agentnetwork/access_log_handler.go new file mode 100644 index 000000000..0a7d5fd58 --- /dev/null +++ b/management/server/http/handlers/agentnetwork/access_log_handler.go @@ -0,0 +1,91 @@ +package agentnetwork + +import ( + "net/http" + "time" + + "github.com/gorilla/mux" + + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + nbcontext "github.com/netbirdio/netbird/management/server/context" + "github.com/netbirdio/netbird/shared/management/http/api" + "github.com/netbirdio/netbird/shared/management/http/util" +) + +// addAccessLogEndpoints registers the read-only, server-side-filtered +// agent-network access-log listing and the aggregated usage overview. +func (h *handler) addAccessLogEndpoints(router *mux.Router) { + router.HandleFunc("/agent-network/access-logs", h.listAccessLogs).Methods("GET", "OPTIONS") + router.HandleFunc("/agent-network/usage/overview", h.getUsageOverview).Methods("GET", "OPTIONS") +} + +func (h *handler) getUsageOverview(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + // Reuse the access-log filter for the shared date/user/group/provider/model + // params; pagination/sort/search are irrelevant for an aggregate. + var filter types.AgentNetworkAccessLogFilter + if err := filter.ParseFromRequest(r); err != nil { + util.WriteError(r.Context(), err, w) + return + } + // Bound the aggregation window so an unbounded or over-wide query can't load + // an account's entire usage history into memory. + filter.ApplyUsageOverviewBounds(time.Now()) + granularity := types.ParseUsageGranularity(r.URL.Query().Get("granularity")) + + buckets, err := h.manager.GetUsageOverview(r.Context(), userAuth.AccountId, userAuth.UserId, filter, granularity) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + out := make([]api.AgentNetworkUsageBucket, 0, len(buckets)) + for _, b := range buckets { + out = append(out, b.ToAPIResponse()) + } + util.WriteJSONObject(r.Context(), w, out) +} + +func (h *handler) listAccessLogs(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + var filter types.AgentNetworkAccessLogFilter + if err := filter.ParseFromRequest(r); err != nil { + util.WriteError(r.Context(), err, w) + return + } + + rows, total, err := h.manager.ListAccessLogs(r.Context(), userAuth.AccountId, userAuth.UserId, filter) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + data := make([]api.AgentNetworkAccessLog, 0, len(rows)) + for _, row := range rows { + data = append(data, row.ToAPIResponse()) + } + + pageSize := filter.GetLimit() + totalPages := 0 + if pageSize > 0 { + totalPages = int((total + int64(pageSize) - 1) / int64(pageSize)) + } + + util.WriteJSONObject(r.Context(), w, api.AgentNetworkAccessLogsResponse{ + Data: data, + Page: filter.Page, + PageSize: pageSize, + TotalRecords: int(total), + TotalPages: totalPages, + }) +} diff --git a/management/server/http/handlers/agentnetwork/budget_handler.go b/management/server/http/handlers/agentnetwork/budget_handler.go new file mode 100644 index 000000000..9dd6ea0ed --- /dev/null +++ b/management/server/http/handlers/agentnetwork/budget_handler.go @@ -0,0 +1,172 @@ +package agentnetwork + +import ( + "encoding/json" + "net/http" + "strings" + + "github.com/gorilla/mux" + + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + nbcontext "github.com/netbirdio/netbird/management/server/context" + "github.com/netbirdio/netbird/shared/management/http/api" + "github.com/netbirdio/netbird/shared/management/http/util" + "github.com/netbirdio/netbird/shared/management/status" +) + +// addBudgetRuleEndpoints registers the account-level budget rule routes. +func (h *handler) addBudgetRuleEndpoints(router *mux.Router) { + router.HandleFunc("/agent-network/budget-rules", h.getAllBudgetRules).Methods("GET", "OPTIONS") + router.HandleFunc("/agent-network/budget-rules", h.createBudgetRule).Methods("POST", "OPTIONS") + router.HandleFunc("/agent-network/budget-rules/{ruleId}", h.getBudgetRule).Methods("GET", "OPTIONS") + router.HandleFunc("/agent-network/budget-rules/{ruleId}", h.updateBudgetRule).Methods("PUT", "OPTIONS") + router.HandleFunc("/agent-network/budget-rules/{ruleId}", h.deleteBudgetRule).Methods("DELETE", "OPTIONS") +} + +func (h *handler) getAllBudgetRules(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + rules, err := h.manager.GetAllBudgetRules(r.Context(), userAuth.AccountId, userAuth.UserId) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + out := make([]*api.AgentNetworkBudgetRule, 0, len(rules)) + for _, rule := range rules { + out = append(out, rule.ToAPIResponse()) + } + util.WriteJSONObject(r.Context(), w, out) +} + +func (h *handler) getBudgetRule(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + ruleID := mux.Vars(r)["ruleId"] + if ruleID == "" { + util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "budget rule ID is required"), w) + return + } + + rule, err := h.manager.GetBudgetRule(r.Context(), userAuth.AccountId, userAuth.UserId, ruleID) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + util.WriteJSONObject(r.Context(), w, rule.ToAPIResponse()) +} + +func (h *handler) createBudgetRule(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + var req api.AgentNetworkBudgetRuleRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + util.WriteErrorResponse("couldn't parse JSON request", http.StatusBadRequest, w) + return + } + + if err := validateBudgetRule(&req); err != nil { + util.WriteError(r.Context(), err, w) + return + } + + rule := types.NewAccountBudgetRule(userAuth.AccountId) + rule.FromAPIRequest(&req) + + created, err := h.manager.CreateBudgetRule(r.Context(), userAuth.UserId, rule) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + util.WriteJSONObject(r.Context(), w, created.ToAPIResponse()) +} + +func (h *handler) updateBudgetRule(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + ruleID := mux.Vars(r)["ruleId"] + if ruleID == "" { + util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "budget rule ID is required"), w) + return + } + + var req api.AgentNetworkBudgetRuleRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + util.WriteErrorResponse("couldn't parse JSON request", http.StatusBadRequest, w) + return + } + + if err := validateBudgetRule(&req); err != nil { + util.WriteError(r.Context(), err, w) + return + } + + rule := &types.AccountBudgetRule{ID: ruleID, AccountID: userAuth.AccountId} + rule.FromAPIRequest(&req) + + updated, err := h.manager.UpdateBudgetRule(r.Context(), userAuth.UserId, rule) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + util.WriteJSONObject(r.Context(), w, updated.ToAPIResponse()) +} + +func (h *handler) deleteBudgetRule(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + ruleID := mux.Vars(r)["ruleId"] + if ruleID == "" { + util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "budget rule ID is required"), w) + return + } + + if err := h.manager.DeleteBudgetRule(r.Context(), userAuth.AccountId, userAuth.UserId, ruleID); err != nil { + util.WriteError(r.Context(), err, w) + return + } + util.WriteJSONObject(r.Context(), w, util.EmptyObject{}) +} + +// validateBudgetRule rejects malformed budget rules. It reuses the policy limit +// validation since the cap shape is identical, and rejects empty target entries. +func validateBudgetRule(req *api.AgentNetworkBudgetRuleRequest) error { + if strings.TrimSpace(req.Name) == "" { + return status.Errorf(status.InvalidArgument, "name is required") + } + if req.TargetGroups != nil { + for _, id := range *req.TargetGroups { + if strings.TrimSpace(id) == "" { + return status.Errorf(status.InvalidArgument, "target_groups must not contain empty entries") + } + } + } + if req.TargetUsers != nil { + for _, id := range *req.TargetUsers { + if strings.TrimSpace(id) == "" { + return status.Errorf(status.InvalidArgument, "target_users must not contain empty entries") + } + } + } + return validatePolicyLimits(req.Limits) +} diff --git a/management/server/http/handlers/agentnetwork/budget_handler_test.go b/management/server/http/handlers/agentnetwork/budget_handler_test.go new file mode 100644 index 000000000..f3c09ee96 --- /dev/null +++ b/management/server/http/handlers/agentnetwork/budget_handler_test.go @@ -0,0 +1,131 @@ +package agentnetwork + +import ( + "context" + "encoding/json" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + agentNetworkTypes "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// TestBudgetRuleHandler_RoundTrip seeds a budget rule via the store and asserts +// the GET wire shape carries targets and the reused PolicyLimits cap shape. The +// create/update/delete success paths go through accountManager.StoreEvent which +// this fixture doesn't wire — they are covered by the manager-level no-mock +// test (TestAgentNetwork_BudgetRuleCRUD_RealManager). +func TestBudgetRuleHandler_RoundTrip(t *testing.T) { + f := newAgentNetworkHandlerFixture(t) + + rule := &agentNetworkTypes.AccountBudgetRule{ + ID: "ainbud_test", + AccountID: testAccountID, + Name: "org-monthly", + Enabled: true, + TargetGroups: []string{"grp-eng"}, + TargetUsers: []string{"user-alice"}, + Limits: agentNetworkTypes.PolicyLimits{ + TokenLimit: agentNetworkTypes.PolicyTokenLimit{Enabled: true, GroupCap: 100000, UserCap: 10000, WindowSeconds: 2_592_000}, + BudgetLimit: agentNetworkTypes.PolicyBudgetLimit{Enabled: true, GroupCapUsd: 500, WindowSeconds: 2_592_000}, + }, + } + require.NoError(t, f.store.SaveAgentNetworkBudgetRule(context.Background(), rule)) + + rec := f.do(t, http.MethodGet, "/agent-network/budget-rules/"+rule.ID, "") + require.Equal(t, http.StatusOK, rec.Code, "GET must succeed: %s", rec.Body.String()) + + var got api.AgentNetworkBudgetRule + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + assert.Equal(t, "org-monthly", got.Name, "name must round-trip") + assert.Equal(t, []string{"grp-eng"}, got.TargetGroups, "target groups must round-trip") + assert.Equal(t, []string{"user-alice"}, got.TargetUsers, "target users must round-trip") + assert.Equal(t, int64(100000), got.Limits.TokenLimit.GroupCap, "token group cap must round-trip") + assert.Equal(t, int64(2_592_000), got.Limits.BudgetLimit.WindowSeconds, "budget window must round-trip") +} + +// TestBudgetRuleHandler_ListReturnsArray asserts the list endpoint returns a +// JSON array (never null) for an account with no rules. +func TestBudgetRuleHandler_ListReturnsArray(t *testing.T) { + f := newAgentNetworkHandlerFixture(t) + + rec := f.do(t, http.MethodGet, "/agent-network/budget-rules", "") + require.Equal(t, http.StatusOK, rec.Code, "GET must succeed: %s", rec.Body.String()) + assert.Equal(t, "[]", trimSpace(rec.Body.String()), "empty account must return an empty array, not null") +} + +// TestBudgetRuleHandler_RejectsMissingName covers the validation path (which +// runs before the manager call, so it works without a wired accountManager). +func TestBudgetRuleHandler_RejectsMissingName(t *testing.T) { + f := newAgentNetworkHandlerFixture(t) + + body := `{ + "name": "", + "limits": { + "token_limit": {"enabled": false, "group_cap": 0, "user_cap": 0, "window_seconds": 0}, + "budget_limit": {"enabled": false, "group_cap_usd": 0, "user_cap_usd": 0, "window_seconds": 0} + } + }` + rec := f.do(t, http.MethodPost, "/agent-network/budget-rules", body) + assert.Equal(t, http.StatusUnprocessableEntity, rec.Code, + "missing name must be rejected as a validation error (not a route/auth 4xx): got %d body=%s", rec.Code, rec.Body.String()) + assert.Contains(t, rec.Body.String(), "name", + "rejection body must name the offending field, proving the validation path: %s", rec.Body.String()) +} + +// TestBudgetRuleHandler_RejectsSubMinuteWindow proves budget rules reuse the +// policy-limit validation (enabled limit needs window >= 60s). +func TestBudgetRuleHandler_RejectsSubMinuteWindow(t *testing.T) { + f := newAgentNetworkHandlerFixture(t) + + body := `{ + "name": "bad-window", + "limits": { + "token_limit": {"enabled": true, "group_cap": 1000, "user_cap": 0, "window_seconds": 30}, + "budget_limit": {"enabled": false, "group_cap_usd": 0, "user_cap_usd": 0, "window_seconds": 0} + } + }` + rec := f.do(t, http.MethodPost, "/agent-network/budget-rules", body) + assert.Equal(t, http.StatusUnprocessableEntity, rec.Code, + "sub-minute window must be rejected as a validation error (not a route/auth 4xx): got %d body=%s", rec.Code, rec.Body.String()) + assert.Contains(t, rec.Body.String(), "window_seconds", + "rejection body must name the offending window_seconds field, proving the validation path: %s", rec.Body.String()) +} + +// TestSettingsHandler_GetExposesCollectionToggles asserts the GET settings wire +// shape carries the account-level collection toggles after a store seed. +func TestSettingsHandler_GetExposesCollectionToggles(t *testing.T) { + f := newAgentNetworkHandlerFixture(t) + + require.NoError(t, f.store.SaveAgentNetworkSettings(context.Background(), &agentNetworkTypes.Settings{ + AccountID: testAccountID, + Cluster: "eu.proxy.netbird.io", + Subdomain: "violet", + EnableLogCollection: true, + EnablePromptCollection: true, + RedactPii: false, + })) + + rec := f.do(t, http.MethodGet, "/agent-network/settings", "") + require.Equal(t, http.StatusOK, rec.Code, "GET must succeed: %s", rec.Body.String()) + + var got api.AgentNetworkSettings + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + assert.True(t, got.EnableLogCollection, "log collection toggle must surface on the wire") + assert.True(t, got.EnablePromptCollection, "prompt collection toggle must surface on the wire") + assert.False(t, got.RedactPii, "redact toggle must surface its false value") + assert.Equal(t, "violet.eu.proxy.netbird.io", got.Endpoint, "endpoint stays computed from immutable cluster+subdomain") +} + +func trimSpace(s string) string { + for len(s) > 0 && (s[len(s)-1] == '\n' || s[len(s)-1] == ' ' || s[len(s)-1] == '\t' || s[len(s)-1] == '\r') { + s = s[:len(s)-1] + } + for len(s) > 0 && (s[0] == '\n' || s[0] == ' ' || s[0] == '\t' || s[0] == '\r') { + s = s[1:] + } + return s +} diff --git a/management/server/http/handlers/agentnetwork/consumption_handler.go b/management/server/http/handlers/agentnetwork/consumption_handler.go new file mode 100644 index 000000000..09a9d6b9a --- /dev/null +++ b/management/server/http/handlers/agentnetwork/consumption_handler.go @@ -0,0 +1,53 @@ +package agentnetwork + +import ( + "net/http" + + "github.com/gorilla/mux" + + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + nbcontext "github.com/netbirdio/netbird/management/server/context" + "github.com/netbirdio/netbird/shared/management/http/api" + "github.com/netbirdio/netbird/shared/management/http/util" +) + +// addConsumptionEndpoints registers the read-only Agent Network +// consumption listing — backs the dashboard's basic counter view. +func (h *handler) addConsumptionEndpoints(router *mux.Router) { + router.HandleFunc("/agent-network/consumption", h.listConsumption).Methods("GET", "OPTIONS") +} + +func (h *handler) listConsumption(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + rows, err := h.manager.ListConsumption(r.Context(), userAuth.AccountId, userAuth.UserId) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + out := make([]api.AgentNetworkConsumption, 0, len(rows)) + for _, row := range rows { + out = append(out, consumptionToAPI(row)) + } + util.WriteJSONObject(r.Context(), w, out) +} + +func consumptionToAPI(c *types.Consumption) api.AgentNetworkConsumption { + windowStart := c.WindowStartUTC + updatedAt := c.UpdatedAt + return api.AgentNetworkConsumption{ + DimensionKind: api.AgentNetworkConsumptionDimensionKind(c.DimensionKind), + DimensionId: c.DimensionID, + WindowSeconds: c.WindowSeconds, + WindowStartUtc: windowStart, + TokensInput: c.TokensInput, + TokensOutput: c.TokensOutput, + CostUsd: c.CostUSD, + UpdatedAt: &updatedAt, + } +} diff --git a/management/server/http/handlers/agentnetwork/guardrails_handler.go b/management/server/http/handlers/agentnetwork/guardrails_handler.go new file mode 100644 index 000000000..b98bb01f5 --- /dev/null +++ b/management/server/http/handlers/agentnetwork/guardrails_handler.go @@ -0,0 +1,171 @@ +package agentnetwork + +import ( + "encoding/json" + "net/http" + "strings" + + "github.com/gorilla/mux" + + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + nbcontext "github.com/netbirdio/netbird/management/server/context" + "github.com/netbirdio/netbird/shared/management/http/api" + "github.com/netbirdio/netbird/shared/management/http/util" + "github.com/netbirdio/netbird/shared/management/status" +) + +// addGuardrailEndpoints registers all Agent Network guardrail routes. +func (h *handler) addGuardrailEndpoints(router *mux.Router) { + router.HandleFunc("/agent-network/guardrails", h.getAllGuardrails).Methods("GET", "OPTIONS") + router.HandleFunc("/agent-network/guardrails", h.createGuardrail).Methods("POST", "OPTIONS") + router.HandleFunc("/agent-network/guardrails/{guardrailId}", h.getGuardrail).Methods("GET", "OPTIONS") + router.HandleFunc("/agent-network/guardrails/{guardrailId}", h.updateGuardrail).Methods("PUT", "OPTIONS") + router.HandleFunc("/agent-network/guardrails/{guardrailId}", h.deleteGuardrail).Methods("DELETE", "OPTIONS") +} + +func (h *handler) getAllGuardrails(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + guardrails, err := h.manager.GetAllGuardrails(r.Context(), userAuth.AccountId, userAuth.UserId) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + out := make([]*api.AgentNetworkGuardrail, 0, len(guardrails)) + for _, g := range guardrails { + out = append(out, g.ToAPIResponse()) + } + util.WriteJSONObject(r.Context(), w, out) +} + +func (h *handler) getGuardrail(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + guardrailID := mux.Vars(r)["guardrailId"] + if guardrailID == "" { + util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "guardrail ID is required"), w) + return + } + + guardrail, err := h.manager.GetGuardrail(r.Context(), userAuth.AccountId, userAuth.UserId, guardrailID) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + util.WriteJSONObject(r.Context(), w, guardrail.ToAPIResponse()) +} + +func (h *handler) createGuardrail(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + var req api.AgentNetworkGuardrailRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + util.WriteErrorResponse("couldn't parse JSON request", http.StatusBadRequest, w) + return + } + + if err := validateGuardrail(&req); err != nil { + util.WriteError(r.Context(), err, w) + return + } + + guardrail := types.NewGuardrail(userAuth.AccountId) + guardrail.FromAPIRequest(&req) + + created, err := h.manager.CreateGuardrail(r.Context(), userAuth.UserId, guardrail) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + util.WriteJSONObject(r.Context(), w, created.ToAPIResponse()) +} + +func (h *handler) updateGuardrail(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + guardrailID := mux.Vars(r)["guardrailId"] + if guardrailID == "" { + util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "guardrail ID is required"), w) + return + } + + var req api.AgentNetworkGuardrailRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + util.WriteErrorResponse("couldn't parse JSON request", http.StatusBadRequest, w) + return + } + + if err := validateGuardrail(&req); err != nil { + util.WriteError(r.Context(), err, w) + return + } + + guardrail := &types.Guardrail{ + ID: guardrailID, + AccountID: userAuth.AccountId, + } + guardrail.FromAPIRequest(&req) + + updated, err := h.manager.UpdateGuardrail(r.Context(), userAuth.UserId, guardrail) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + util.WriteJSONObject(r.Context(), w, updated.ToAPIResponse()) +} + +func (h *handler) deleteGuardrail(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + guardrailID := mux.Vars(r)["guardrailId"] + if guardrailID == "" { + util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "guardrail ID is required"), w) + return + } + + if err := h.manager.DeleteGuardrail(r.Context(), userAuth.AccountId, userAuth.UserId, guardrailID); err != nil { + util.WriteError(r.Context(), err, w) + return + } + + util.WriteJSONObject(r.Context(), w, util.EmptyObject{}) +} + +func validateGuardrail(req *api.AgentNetworkGuardrailRequest) error { + if strings.TrimSpace(req.Name) == "" { + return status.Errorf(status.InvalidArgument, "name is required") + } + + c := req.Checks + if c.ModelAllowlist.Enabled { + for _, id := range c.ModelAllowlist.Models { + if strings.TrimSpace(id) == "" { + return status.Errorf(status.InvalidArgument, "model_allowlist.models must not contain empty entries") + } + } + } + return nil +} diff --git a/management/server/http/handlers/agentnetwork/handlers_test.go b/management/server/http/handlers/agentnetwork/handlers_test.go new file mode 100644 index 000000000..51b11aa88 --- /dev/null +++ b/management/server/http/handlers/agentnetwork/handlers_test.go @@ -0,0 +1,256 @@ +package agentnetwork + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "runtime" + "strings" + "testing" + + "github.com/golang/mock/gomock" + "github.com/gorilla/mux" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/management/server/agentnetwork" + agentNetworkTypes "github.com/netbirdio/netbird/management/server/agentnetwork/types" + nbcontext "github.com/netbirdio/netbird/management/server/context" + "github.com/netbirdio/netbird/management/server/permissions" + "github.com/netbirdio/netbird/management/server/store" + nbtypes "github.com/netbirdio/netbird/management/server/types" + "github.com/netbirdio/netbird/shared/auth" + "github.com/netbirdio/netbird/shared/management/http/api" +) + +const ( + testAccountID = "acc-1" + testUserID = "user-bob" +) + +// agentNetworkHandlerFixture builds a real agentnetwork.Manager with +// a sqlite store and an always-allow permissions mock, then exposes +// the HTTP handlers via a gorilla router. Tests issue requests +// through httptest and assert on the wire shape — the same path the +// dashboard exercises. +type agentNetworkHandlerFixture struct { + store store.Store + manager agentnetwork.Manager + router *mux.Router +} + +func newAgentNetworkHandlerFixture(t *testing.T) *agentNetworkHandlerFixture { + t.Helper() + if runtime.GOOS == "windows" { + t.Skip("sqlite store not properly supported on Windows yet") + } + t.Setenv("NETBIRD_STORE_ENGINE", string(nbtypes.SqliteStoreEngine)) + + st, cleanUp, err := store.NewTestStoreFromSQL(context.Background(), "", t.TempDir()) + require.NoError(t, err) + t.Cleanup(cleanUp) + + ctrl := gomock.NewController(t) + perms := permissions.NewMockManager(ctrl) + // Always-allow: the handler tests are about wire shape, not + // authz. Authz is covered by the manager's own tests. + perms.EXPECT(). + ValidateUserPermissions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(true, context.Background(), nil). + AnyTimes() + + manager := agentnetwork.NewManager(st, perms, nil, nil) + h := &handler{manager: manager} + + router := mux.NewRouter() + h.addPolicyEndpoints(router) + h.addConsumptionEndpoints(router) + h.addBudgetRuleEndpoints(router) + h.addSettingsEndpoints(router) + + return &agentNetworkHandlerFixture{ + store: st, + manager: manager, + router: router, + } +} + +func (f *agentNetworkHandlerFixture) do(t *testing.T, method, path, body string) *httptest.ResponseRecorder { + t.Helper() + var reader io.Reader + if body != "" { + reader = strings.NewReader(body) + } + req := httptest.NewRequest(method, path, reader) + if body != "" { + req.Header.Set("Content-Type", "application/json") + } + req = nbcontext.SetUserAuthInRequest(req, auth.UserAuth{ + UserId: testUserID, + AccountId: testAccountID, + }) + rec := httptest.NewRecorder() + f.router.ServeHTTP(rec, req) + return rec +} + +// seedProvider persists a minimal provider record so policy create +// passes the manager's destination_provider_ids existence check. +func (f *agentNetworkHandlerFixture) seedProvider(t *testing.T, id string) { + t.Helper() + require.NoError(t, f.store.SaveAgentNetworkProvider(context.Background(), &agentNetworkTypes.Provider{ + ID: id, + AccountID: testAccountID, + ProviderID: "openai_api", + Name: "test-" + id, + UpstreamURL: "https://api.openai.com", + APIKey: "sk-test", + Enabled: true, + SessionPrivateKey: "test-priv-key", + SessionPublicKey: "test-pub-key", + })) +} + +// TestPolicyHandler_WindowSecondsRoundTrip ports bash 10 to Go: +// assert that a policy with window_seconds on both Token + Budget +// halves round-trips through GET unchanged AND that legacy +// window_hours / window_days are absent from the JSON response. We +// seed the policy directly via the store rather than POST-ing +// because the create path goes through the manager's +// accountManager.StoreEvent which we don't wire in this fixture; the +// on-wire shape is what matters here, and the POST validation path +// is covered separately by the RejectsSubMinuteWindow test. +func TestPolicyHandler_WindowSecondsRoundTrip(t *testing.T) { + f := newAgentNetworkHandlerFixture(t) + + policy := &agentNetworkTypes.Policy{ + ID: "ainpol_test", + AccountID: testAccountID, + Name: "round-trip", + Enabled: true, + SourceGroups: []string{"grp-engineers"}, + DestinationProviderIDs: []string{"prov-1"}, + Limits: agentNetworkTypes.PolicyLimits{ + TokenLimit: agentNetworkTypes.PolicyTokenLimit{Enabled: true, GroupCap: 10000, UserCap: 5000, WindowSeconds: 86_400}, + BudgetLimit: agentNetworkTypes.PolicyBudgetLimit{Enabled: true, GroupCapUsd: 10.0, UserCapUsd: 2.5, WindowSeconds: 2_592_000}, + }, + } + require.NoError(t, f.store.SaveAgentNetworkPolicy(context.Background(), policy)) + + rec := f.do(t, http.MethodGet, "/agent-network/policies/"+policy.ID, "") + require.Equal(t, http.StatusOK, rec.Code, "GET must succeed: %s", rec.Body.String()) + + var got api.AgentNetworkPolicy + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + assert.Equal(t, int64(86_400), got.Limits.TokenLimit.WindowSeconds, "token_limit.window_seconds must round-trip") + assert.Equal(t, int64(2_592_000), got.Limits.BudgetLimit.WindowSeconds, "budget_limit.window_seconds must round-trip") + + // Legacy field names must NOT appear in the response — would + // signal that the management server is still emitting the old + // shape and would fool a v1 dashboard into rendering days/hours. + assert.NotContains(t, rec.Body.String(), "window_hours", + "legacy window_hours field must be absent from the on-wire response") + assert.NotContains(t, rec.Body.String(), "window_days", + "legacy window_days field must be absent from the on-wire response") +} + +// TestPolicyHandler_RejectsSubMinuteWindow ports bash 20 to Go: an +// enabled limit with window_seconds < 60 must surface as a 4xx +// because anything finer than per-minute produces an untenable +// volume of consumption rows for a feature whose value comes from +// per-window cap enforcement. +func TestPolicyHandler_RejectsSubMinuteWindow(t *testing.T) { + f := newAgentNetworkHandlerFixture(t) + f.seedProvider(t, "prov-1") + + body := `{ + "name": "sub-minute-window", + "enabled": true, + "source_groups": ["grp-engineers"], + "destination_provider_ids": ["prov-1"], + "guardrail_ids": [], + "limits": { + "token_limit": {"enabled": true, "group_cap": 10000, "user_cap": 5000, "window_seconds": 30}, + "budget_limit": {"enabled": false, "group_cap_usd": 0, "user_cap_usd": 0, "window_seconds": 0} + } + }` + rec := f.do(t, http.MethodPost, "/agent-network/policies", body) + // 422 specifically (InvalidArgument) proves the window-validation path — + // a route miss would be 404 and an auth failure 403, so a generic 4xx + // would let those false-pass. + assert.Equal(t, http.StatusUnprocessableEntity, rec.Code, + "enabled token_limit with window_seconds<60 must be rejected as a validation error: got %d body=%s", rec.Code, rec.Body.String()) + assert.Contains(t, rec.Body.String(), "window_seconds", + "rejection body must name the offending window_seconds field, proving it's the validation path: %s", rec.Body.String()) +} + +// TestConsumptionHandler_EmptyAccountReturnsArray ports bash 30 to +// Go: GET /agent-network/consumption on a clean account always +// returns a JSON array (possibly empty), never a 404 / 500. The +// dashboard depends on this shape to render its empty state. +func TestConsumptionHandler_EmptyAccountReturnsArray(t *testing.T) { + f := newAgentNetworkHandlerFixture(t) + + rec := f.do(t, http.MethodGet, "/agent-network/consumption", "") + require.Equal(t, http.StatusOK, rec.Code) + + var rows []api.AgentNetworkConsumption + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &rows), + "response must always be a JSON array — even when empty: %s", rec.Body.String()) + assert.Empty(t, rows) +} + +// TestConsumptionHandler_PopulatedAccountListsRows mirrors the +// /consumption read after a few RecordConsumption calls. Validates +// the wire shape carries every field the dashboard reads (dim_kind, +// dim_id, window_seconds, window_start_utc, tokens, cost_usd) and +// rows are ordered window-newest-first. +func TestConsumptionHandler_PopulatedAccountListsRows(t *testing.T) { + f := newAgentNetworkHandlerFixture(t) + + require.NoError(t, f.manager.RecordConsumption( + context.Background(), testAccountID, + agentNetworkTypes.DimensionGroup, "grp-engineers", + 86_400, 100, 50, 0.0125, + )) + require.NoError(t, f.manager.RecordConsumption( + context.Background(), testAccountID, + agentNetworkTypes.DimensionUser, testUserID, + 86_400, 100, 50, 0.0125, + )) + + rec := f.do(t, http.MethodGet, "/agent-network/consumption", "") + require.Equal(t, http.StatusOK, rec.Code) + + var rows []api.AgentNetworkConsumption + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &rows)) + require.Len(t, rows, 2, "two RecordConsumption calls must yield two rows") + + // Index by dim_kind so we can assert the full wire shape of each row, + // including the dimension id and the aligned window start the dashboard + // keys on. Both rows share totals and window. + byKind := make(map[string]api.AgentNetworkConsumption, len(rows)) + for _, row := range rows { + assert.Equal(t, int64(100), row.TokensInput) + assert.Equal(t, int64(50), row.TokensOutput) + assert.InDelta(t, 0.0125, row.CostUsd, 1e-9) + assert.Equal(t, int64(86_400), row.WindowSeconds) + assert.False(t, row.WindowStartUtc.IsZero(), "window_start_utc must be set on every row") + byKind[string(row.DimensionKind)] = row + } + + groupRow, ok := byKind["group"] + require.True(t, ok, "group dimension must surface") + assert.Equal(t, "grp-engineers", groupRow.DimensionId, "group row must carry the source group id as dimension_id") + + userRow, ok := byKind["user"] + require.True(t, ok, "user dimension must surface") + assert.Equal(t, testUserID, userRow.DimensionId, "user row must carry the user id as dimension_id") + + // Both rows fall in the same aligned window (same length, recorded + // together), so window_start_utc must match across them. + assert.Equal(t, groupRow.WindowStartUtc, userRow.WindowStartUtc, + "rows recorded in the same window must share the aligned window_start_utc") +} diff --git a/management/server/http/handlers/agentnetwork/policies_handler.go b/management/server/http/handlers/agentnetwork/policies_handler.go new file mode 100644 index 000000000..7e5a3e8cb --- /dev/null +++ b/management/server/http/handlers/agentnetwork/policies_handler.go @@ -0,0 +1,228 @@ +package agentnetwork + +import ( + "encoding/json" + "net/http" + "strings" + + "github.com/gorilla/mux" + + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + nbcontext "github.com/netbirdio/netbird/management/server/context" + "github.com/netbirdio/netbird/shared/management/http/api" + "github.com/netbirdio/netbird/shared/management/http/util" + "github.com/netbirdio/netbird/shared/management/status" +) + +// minWindowSeconds is the floor enforced on enabled token / budget +// limit windows. One minute is short enough for fine-grained burst +// control without producing untenable consumption-row volume at scale. +const minWindowSeconds int64 = 60 + +// addPolicyEndpoints registers all Agent Network policy routes on the +// shared handler. +func (h *handler) addPolicyEndpoints(router *mux.Router) { + router.HandleFunc("/agent-network/policies", h.getAllPolicies).Methods("GET", "OPTIONS") + router.HandleFunc("/agent-network/policies", h.createPolicy).Methods("POST", "OPTIONS") + router.HandleFunc("/agent-network/policies/{policyId}", h.getPolicy).Methods("GET", "OPTIONS") + router.HandleFunc("/agent-network/policies/{policyId}", h.updatePolicy).Methods("PUT", "OPTIONS") + router.HandleFunc("/agent-network/policies/{policyId}", h.deletePolicy).Methods("DELETE", "OPTIONS") +} + +func (h *handler) getAllPolicies(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + policies, err := h.manager.GetAllPolicies(r.Context(), userAuth.AccountId, userAuth.UserId) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + out := make([]*api.AgentNetworkPolicy, 0, len(policies)) + for _, p := range policies { + out = append(out, p.ToAPIResponse()) + } + util.WriteJSONObject(r.Context(), w, out) +} + +func (h *handler) getPolicy(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + policyID := mux.Vars(r)["policyId"] + if policyID == "" { + util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "policy ID is required"), w) + return + } + + policy, err := h.manager.GetPolicy(r.Context(), userAuth.AccountId, userAuth.UserId, policyID) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + util.WriteJSONObject(r.Context(), w, policy.ToAPIResponse()) +} + +func (h *handler) createPolicy(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + var req api.AgentNetworkPolicyRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + util.WriteErrorResponse("couldn't parse JSON request", http.StatusBadRequest, w) + return + } + + if err := validatePolicy(&req); err != nil { + util.WriteError(r.Context(), err, w) + return + } + + policy := types.NewPolicy(userAuth.AccountId) + policy.FromAPIRequest(&req) + + created, err := h.manager.CreatePolicy(r.Context(), userAuth.UserId, policy) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + util.WriteJSONObject(r.Context(), w, created.ToAPIResponse()) +} + +func (h *handler) updatePolicy(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + policyID := mux.Vars(r)["policyId"] + if policyID == "" { + util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "policy ID is required"), w) + return + } + + var req api.AgentNetworkPolicyRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + util.WriteErrorResponse("couldn't parse JSON request", http.StatusBadRequest, w) + return + } + + if err := validatePolicy(&req); err != nil { + util.WriteError(r.Context(), err, w) + return + } + + policy := &types.Policy{ + ID: policyID, + AccountID: userAuth.AccountId, + } + policy.FromAPIRequest(&req) + + updated, err := h.manager.UpdatePolicy(r.Context(), userAuth.UserId, policy) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + util.WriteJSONObject(r.Context(), w, updated.ToAPIResponse()) +} + +func (h *handler) deletePolicy(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + policyID := mux.Vars(r)["policyId"] + if policyID == "" { + util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "policy ID is required"), w) + return + } + + if err := h.manager.DeletePolicy(r.Context(), userAuth.AccountId, userAuth.UserId, policyID); err != nil { + util.WriteError(r.Context(), err, w) + return + } + + util.WriteJSONObject(r.Context(), w, util.EmptyObject{}) +} + +func validatePolicy(req *api.AgentNetworkPolicyRequest) error { + if strings.TrimSpace(req.Name) == "" { + return status.Errorf(status.InvalidArgument, "name is required") + } + if len(req.SourceGroups) == 0 { + return status.Errorf(status.InvalidArgument, "source_groups must contain at least one group id") + } + for _, id := range req.SourceGroups { + if strings.TrimSpace(id) == "" { + return status.Errorf(status.InvalidArgument, "source_groups must not contain empty entries") + } + } + if len(req.DestinationProviderIds) == 0 { + return status.Errorf(status.InvalidArgument, "destination_provider_ids must contain at least one provider id") + } + for _, id := range req.DestinationProviderIds { + if strings.TrimSpace(id) == "" { + return status.Errorf(status.InvalidArgument, "destination_provider_ids must not contain empty entries") + } + } + if req.GuardrailIds != nil { + for _, id := range *req.GuardrailIds { + if strings.TrimSpace(id) == "" { + return status.Errorf(status.InvalidArgument, "guardrail_ids must not contain empty entries") + } + } + } + if req.Limits != nil { + if err := validatePolicyLimits(*req.Limits); err != nil { + return err + } + } + return nil +} + +func validatePolicyLimits(l api.AgentNetworkPolicyLimits) error { + if l.TokenLimit.Enabled { + if l.TokenLimit.WindowSeconds < minWindowSeconds { + return status.Errorf(status.InvalidArgument, "limits.token_limit.window_seconds must be at least %d (one minute) when enabled", minWindowSeconds) + } + if l.TokenLimit.GroupCap < 0 { + return status.Errorf(status.InvalidArgument, "limits.token_limit.group_cap must not be negative") + } + if l.TokenLimit.UserCap < 0 { + return status.Errorf(status.InvalidArgument, "limits.token_limit.user_cap must not be negative") + } + if l.TokenLimit.GroupCap == 0 && l.TokenLimit.UserCap == 0 { + return status.Errorf(status.InvalidArgument, "limits.token_limit requires group_cap or user_cap to be greater than zero when enabled") + } + } + if l.BudgetLimit.Enabled { + if l.BudgetLimit.WindowSeconds < minWindowSeconds { + return status.Errorf(status.InvalidArgument, "limits.budget_limit.window_seconds must be at least %d (one minute) when enabled", minWindowSeconds) + } + if l.BudgetLimit.GroupCapUsd < 0 { + return status.Errorf(status.InvalidArgument, "limits.budget_limit.group_cap_usd must not be negative") + } + if l.BudgetLimit.UserCapUsd < 0 { + return status.Errorf(status.InvalidArgument, "limits.budget_limit.user_cap_usd must not be negative") + } + if l.BudgetLimit.GroupCapUsd == 0 && l.BudgetLimit.UserCapUsd == 0 { + return status.Errorf(status.InvalidArgument, "limits.budget_limit requires group_cap_usd or user_cap_usd to be greater than zero when enabled") + } + } + return nil +} diff --git a/management/server/http/handlers/agentnetwork/providers_handler.go b/management/server/http/handlers/agentnetwork/providers_handler.go new file mode 100644 index 000000000..b05dc83aa --- /dev/null +++ b/management/server/http/handlers/agentnetwork/providers_handler.go @@ -0,0 +1,217 @@ +// Package agentnetwork serves the Agent Network HTTP API. +// +// All persistence is delegated to agentnetwork.Manager so this layer only +// translates between the wire format (api.AgentNetworkProvider*) and the +// domain types. +package agentnetwork + +import ( + "encoding/json" + "net/http" + "net/url" + "strings" + + "github.com/gorilla/mux" + + "github.com/netbirdio/netbird/management/server/agentnetwork" + "github.com/netbirdio/netbird/management/server/agentnetwork/catalog" + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + nbcontext "github.com/netbirdio/netbird/management/server/context" + "github.com/netbirdio/netbird/shared/management/http/api" + "github.com/netbirdio/netbird/shared/management/http/util" + "github.com/netbirdio/netbird/shared/management/status" +) + +type handler struct { + manager agentnetwork.Manager +} + +// AddEndpoints registers all Agent Network routes. +func AddEndpoints(manager agentnetwork.Manager, router *mux.Router) { + h := &handler{manager: manager} + router.HandleFunc("/agent-network/catalog/providers", h.getCatalogProviders).Methods("GET", "OPTIONS") + router.HandleFunc("/agent-network/providers", h.getAllProviders).Methods("GET", "OPTIONS") + router.HandleFunc("/agent-network/providers", h.createProvider).Methods("POST", "OPTIONS") + router.HandleFunc("/agent-network/providers/{providerId}", h.getProvider).Methods("GET", "OPTIONS") + router.HandleFunc("/agent-network/providers/{providerId}", h.updateProvider).Methods("PUT", "OPTIONS") + router.HandleFunc("/agent-network/providers/{providerId}", h.deleteProvider).Methods("DELETE", "OPTIONS") + h.addPolicyEndpoints(router) + h.addGuardrailEndpoints(router) + h.addSettingsEndpoints(router) + h.addConsumptionEndpoints(router) + h.addAccessLogEndpoints(router) + h.addBudgetRuleEndpoints(router) +} + +func (h *handler) getCatalogProviders(w http.ResponseWriter, r *http.Request) { + if _, err := nbcontext.GetUserAuthFromContext(r.Context()); err != nil { + util.WriteError(r.Context(), err, w) + return + } + + entries := catalog.All() + out := make([]api.AgentNetworkCatalogProvider, 0, len(entries)) + for _, e := range entries { + out = append(out, e.ToAPIResponse()) + } + util.WriteJSONObject(r.Context(), w, out) +} + +func (h *handler) getAllProviders(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + providers, err := h.manager.GetAllProviders(r.Context(), userAuth.AccountId, userAuth.UserId) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + out := make([]*api.AgentNetworkProvider, 0, len(providers)) + for _, p := range providers { + out = append(out, p.ToAPIResponse()) + } + util.WriteJSONObject(r.Context(), w, out) +} + +func (h *handler) getProvider(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + providerID := mux.Vars(r)["providerId"] + if providerID == "" { + util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "provider ID is required"), w) + return + } + + provider, err := h.manager.GetProvider(r.Context(), userAuth.AccountId, userAuth.UserId, providerID) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + util.WriteJSONObject(r.Context(), w, provider.ToAPIResponse()) +} + +func (h *handler) createProvider(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + var req api.AgentNetworkProviderRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + util.WriteErrorResponse("couldn't parse JSON request", http.StatusBadRequest, w) + return + } + + if err := validate(&req, true); err != nil { + util.WriteError(r.Context(), err, w) + return + } + + provider := types.NewProvider(userAuth.AccountId) + provider.FromAPIRequest(&req) + + bootstrapCluster := "" + if req.BootstrapCluster != nil { + bootstrapCluster = *req.BootstrapCluster + } + + created, err := h.manager.CreateProvider(r.Context(), userAuth.UserId, provider, bootstrapCluster) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + util.WriteJSONObject(r.Context(), w, created.ToAPIResponse()) +} + +func (h *handler) updateProvider(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + providerID := mux.Vars(r)["providerId"] + if providerID == "" { + util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "provider ID is required"), w) + return + } + + var req api.AgentNetworkProviderRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + util.WriteErrorResponse("couldn't parse JSON request", http.StatusBadRequest, w) + return + } + + if err := validate(&req, false); err != nil { + util.WriteError(r.Context(), err, w) + return + } + + provider := &types.Provider{ + ID: providerID, + AccountID: userAuth.AccountId, + } + provider.FromAPIRequest(&req) + + updated, err := h.manager.UpdateProvider(r.Context(), userAuth.UserId, provider) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + util.WriteJSONObject(r.Context(), w, updated.ToAPIResponse()) +} + +func (h *handler) deleteProvider(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + providerID := mux.Vars(r)["providerId"] + if providerID == "" { + util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "provider ID is required"), w) + return + } + + if err := h.manager.DeleteProvider(r.Context(), userAuth.AccountId, userAuth.UserId, providerID); err != nil { + util.WriteError(r.Context(), err, w) + return + } + + util.WriteJSONObject(r.Context(), w, util.EmptyObject{}) +} + +func validate(req *api.AgentNetworkProviderRequest, requireAPIKey bool) error { + if strings.TrimSpace(req.ProviderId) == "" { + return status.Errorf(status.InvalidArgument, "provider_id is required") + } + if !catalog.IsKnown(req.ProviderId) { + return status.Errorf(status.InvalidArgument, "provider_id %q is not a known catalog provider", req.ProviderId) + } + if strings.TrimSpace(req.Name) == "" { + return status.Errorf(status.InvalidArgument, "name is required") + } + if strings.TrimSpace(req.UpstreamUrl) == "" { + return status.Errorf(status.InvalidArgument, "upstream_url is required") + } + u, err := url.Parse(strings.TrimSpace(req.UpstreamUrl)) + if err != nil || u.Host == "" || (u.Scheme != "http" && u.Scheme != "https") { + return status.Errorf(status.InvalidArgument, "upstream_url must be a full http(s) URL") + } + if requireAPIKey && (req.ApiKey == nil || strings.TrimSpace(*req.ApiKey) == "") { + return status.Errorf(status.InvalidArgument, "api_key is required") + } + return nil +} diff --git a/management/server/http/handlers/agentnetwork/settings_handler.go b/management/server/http/handlers/agentnetwork/settings_handler.go new file mode 100644 index 000000000..f28486790 --- /dev/null +++ b/management/server/http/handlers/agentnetwork/settings_handler.go @@ -0,0 +1,74 @@ +package agentnetwork + +import ( + "encoding/json" + "errors" + "net/http" + + "github.com/gorilla/mux" + + "github.com/netbirdio/netbird/management/server/agentnetwork/types" + nbcontext "github.com/netbirdio/netbird/management/server/context" + "github.com/netbirdio/netbird/shared/management/http/api" + "github.com/netbirdio/netbird/shared/management/http/util" + "github.com/netbirdio/netbird/shared/management/status" +) + +// addSettingsEndpoints registers the Agent Network settings routes. The +// settings row is bootstrapped server-side on first provider create; GET reads +// it and PUT updates the mutable collection toggles (cluster/subdomain stay +// immutable). +func (h *handler) addSettingsEndpoints(router *mux.Router) { + router.HandleFunc("/agent-network/settings", h.getSettings).Methods("GET", "OPTIONS") + router.HandleFunc("/agent-network/settings", h.updateSettings).Methods("PUT", "OPTIONS") +} + +// updateSettings applies the collection toggles to the account's settings row. +func (h *handler) updateSettings(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + var req api.AgentNetworkSettingsRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + util.WriteErrorResponse("couldn't parse JSON request", http.StatusBadRequest, w) + return + } + + settings := &types.Settings{AccountID: userAuth.AccountId} + settings.FromAPIRequest(&req) + + updated, err := h.manager.UpdateSettings(r.Context(), userAuth.UserId, settings) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + util.WriteJSONObject(r.Context(), w, updated.ToAPIResponse()) +} + +// getSettings returns the account's agent-network settings. The settings +// row is bootstrapped on first provider create, so freshly-onboarded +// accounts have nothing to read. Rather than 404-ing in that case (which +// the dashboard would have to special-case), return a JSON null with 200 +// so consumers can branch on the body alone. +func (h *handler) getSettings(w http.ResponseWriter, r *http.Request) { + userAuth, err := nbcontext.GetUserAuthFromContext(r.Context()) + if err != nil { + util.WriteError(r.Context(), err, w) + return + } + + settings, err := h.manager.GetSettings(r.Context(), userAuth.AccountId, userAuth.UserId) + if err != nil { + var sErr *status.Error + if errors.As(err, &sErr) && sErr.Type() == status.NotFound { + util.WriteJSONObject(r.Context(), w, nil) + return + } + util.WriteError(r.Context(), err, w) + return + } + util.WriteJSONObject(r.Context(), w, settings.ToAPIResponse()) +} diff --git a/management/server/http/testing/testing_tools/channel/channel.go b/management/server/http/testing/testing_tools/channel/channel.go index 61584a615..8b05b2ddf 100644 --- a/management/server/http/testing/testing_tools/channel/channel.go +++ b/management/server/http/testing/testing_tools/channel/channel.go @@ -137,7 +137,7 @@ func BuildApiBlackBoxWithDBState(t testing_tools.TB, sqlFile string, expectedPee zoneRecordsManager := recordsManager.NewManager(store, am, permissionsManager) apiRouter := mux.NewRouter().PathPrefix("/api").Subrouter() - apiHandler, err := http2.NewAPIHandler(context.Background(), apiRouter, am, networksManager, resourcesManager, routersManager, groupsManager, geoMock, authManagerMock, metrics, permissionsManager, settingsManager, customZonesManager, zoneRecordsManager, networkMapController, nil, serviceManager, nil, nil, nil, nil, nil, nil) + apiHandler, err := http2.NewAPIHandler(context.Background(), apiRouter, am, networksManager, resourcesManager, routersManager, groupsManager, geoMock, authManagerMock, metrics, permissionsManager, settingsManager, customZonesManager, zoneRecordsManager, networkMapController, nil, serviceManager, nil, nil, nil, nil, nil, nil, nil) if err != nil { t.Fatalf("Failed to create API handler: %v", err) } @@ -267,7 +267,7 @@ func BuildApiBlackBoxWithDBStateAndPeerChannel(t testing_tools.TB, sqlFile strin zoneRecordsManager := recordsManager.NewManager(store, am, permissionsManager) apiRouter := mux.NewRouter().PathPrefix("/api").Subrouter() - apiHandler, err := http2.NewAPIHandler(context.Background(), apiRouter, am, networksManager, resourcesManager, routersManager, groupsManager, geoMock, authManagerMock, metrics, permissionsManager, settingsManager, customZonesManager, zoneRecordsManager, networkMapController, nil, serviceManager, nil, nil, nil, nil, nil, nil) + apiHandler, err := http2.NewAPIHandler(context.Background(), apiRouter, am, networksManager, resourcesManager, routersManager, groupsManager, geoMock, authManagerMock, metrics, permissionsManager, settingsManager, customZonesManager, zoneRecordsManager, networkMapController, nil, serviceManager, nil, nil, nil, nil, nil, nil, nil) if err != nil { t.Fatalf("Failed to create API handler: %v", err) } diff --git a/management/server/permissions/modules/module.go b/management/server/permissions/modules/module.go index 93007d4c1..a3a9c554d 100644 --- a/management/server/permissions/modules/module.go +++ b/management/server/permissions/modules/module.go @@ -19,6 +19,7 @@ const ( Pats Module = "pats" IdentityProviders Module = "identity_providers" Services Module = "services" + AgentNetwork Module = "agent_network" ) var All = map[Module]struct{}{ @@ -38,4 +39,5 @@ var All = map[Module]struct{}{ Pats: {}, IdentityProviders: {}, Services: {}, + AgentNetwork: {}, } diff --git a/management/server/store/sql_store.go b/management/server/store/sql_store.go index 8bc4bcd7d..1ff9d92d4 100644 --- a/management/server/store/sql_store.go +++ b/management/server/store/sql_store.go @@ -37,6 +37,7 @@ import ( rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" "github.com/netbirdio/netbird/management/internals/modules/zones" "github.com/netbirdio/netbird/management/internals/modules/zones/records" + agentNetworkTypes "github.com/netbirdio/netbird/management/server/agentnetwork/types" 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" @@ -137,6 +138,10 @@ func NewSqlStore(ctx context.Context, db *gorm.DB, storeEngine types.Engine, met &networkTypes.Network{}, &routerTypes.NetworkRouter{}, &resourceTypes.NetworkResource{}, &types.AccountOnboarding{}, &types.Job{}, &zones.Zone{}, &records.Record{}, &types.UserInviteRecord{}, &rpservice.Service{}, &rpservice.Target{}, &domain.Domain{}, &accesslogs.AccessLogEntry{}, &proxy.Proxy{}, + &agentNetworkTypes.Provider{}, &agentNetworkTypes.Policy{}, &agentNetworkTypes.Guardrail{}, &agentNetworkTypes.Settings{}, + &agentNetworkTypes.Consumption{}, &agentNetworkTypes.AccountBudgetRule{}, + &agentNetworkTypes.AgentNetworkAccessLog{}, &agentNetworkTypes.AgentNetworkAccessLogGroup{}, + &agentNetworkTypes.AgentNetworkUsage{}, &agentNetworkTypes.AgentNetworkUsageGroup{}, ) if err != nil { return nil, fmt.Errorf("auto migratePreAuto: %w", err) @@ -5573,6 +5578,255 @@ func (s *SqlStore) CreateAccessLog(ctx context.Context, logEntry *accesslogs.Acc return nil } +// CreateAgentNetworkAccessLog persists a flattened agent-network access-log +// entry together with its authorising-group child rows in a single +// transaction. +func (s *SqlStore) CreateAgentNetworkAccessLog(ctx context.Context, entry *agentNetworkTypes.AgentNetworkAccessLog, groups []agentNetworkTypes.AgentNetworkAccessLogGroup) error { + err := s.db.Transaction(func(tx *gorm.DB) error { + // Idempotent on the log id / (log_id, group_id) so a proxy resend of the + // same entry can't fail the request. + if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(entry).Error; err != nil { + return err + } + if len(groups) > 0 { + if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&groups).Error; err != nil { + return err + } + } + return nil + }) + if err != nil { + log.WithContext(ctx).WithFields(log.Fields{ + "account_id": entry.AccountID, + "service_id": entry.ServiceID, + "model": entry.Model, + }).Errorf("failed to create agent-network access log entry in store: %v", err) + return status.Errorf(status.Internal, "failed to create agent-network access log entry in store") + } + return nil +} + +// CreateAgentNetworkUsage persists a stripped agent-network usage record +// together with its authorising-group child rows in a single transaction. +func (s *SqlStore) CreateAgentNetworkUsage(ctx context.Context, usage *agentNetworkTypes.AgentNetworkUsage, groups []agentNetworkTypes.AgentNetworkUsageGroup) error { + err := s.db.Transaction(func(tx *gorm.DB) error { + // Idempotent on the usage id / (usage_id, group_id) so a proxy resend of + // the same entry can't fail the request. + if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(usage).Error; err != nil { + return err + } + if len(groups) > 0 { + if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&groups).Error; err != nil { + return err + } + } + return nil + }) + if err != nil { + log.WithContext(ctx).WithFields(log.Fields{ + "account_id": usage.AccountID, + "model": usage.Model, + }).Errorf("failed to create agent-network usage record in store: %v", err) + return status.Errorf(status.Internal, "failed to create agent-network usage record in store") + } + return nil +} + +// DeleteOldAgentNetworkAccessLogs deletes an account's access-log rows (and +// their authorising-group child rows) older than the cutoff. Usage records are +// untouched — they are the long-term aggregate. Returns the number of log rows +// deleted. +func (s *SqlStore) DeleteOldAgentNetworkAccessLogs(ctx context.Context, accountID string, olderThan time.Time) (int64, error) { + var deleted int64 + err := s.db.Transaction(func(tx *gorm.DB) error { + // Remove group child rows for the soon-to-be-deleted logs first. + if err := tx.Exec( + "DELETE FROM agent_network_access_log_group WHERE account_id = ? AND log_id IN (SELECT id FROM agent_network_access_log WHERE account_id = ? AND timestamp < ?)", + accountID, accountID, olderThan, + ).Error; err != nil { + return err + } + res := tx.Where("account_id = ? AND timestamp < ?", accountID, olderThan). + Delete(&agentNetworkTypes.AgentNetworkAccessLog{}) + if res.Error != nil { + return res.Error + } + deleted = res.RowsAffected + return nil + }) + if err != nil { + log.WithContext(ctx).Errorf("failed to delete old agent-network access logs for account %s: %v", accountID, err) + return 0, status.Errorf(status.Internal, "failed to delete old agent-network access logs") + } + return deleted, nil +} + +// GetAgentNetworkUsageRows returns the stripped usage rows for an account that +// match the filter (date / user / group / provider / model). Aggregation into +// time buckets happens in the manager so granularities stay engine-portable. +func (s *SqlStore) GetAgentNetworkUsageRows(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkUsage, error) { + var rows []*agentNetworkTypes.AgentNetworkUsage + + query := s.applyAgentNetworkUsageFilters( + s.db.Where(accountIDCondition, accountID), + filter, + ).Order("timestamp ASC") + + if lockStrength != LockingStrengthNone { + query = query.Clauses(clause.Locking{Strength: string(lockStrength)}) + } + + if err := query.Find(&rows).Error; err != nil { + log.WithContext(ctx).Errorf("failed to get agent-network usage rows from store: %v", err) + return nil, status.Errorf(status.Internal, "failed to get agent-network usage rows from store") + } + return rows, nil +} + +// applyAgentNetworkUsageFilters applies the shared access-log filter's +// date/user/group/provider/model conditions to a usage-table query. Pagination, +// sort and free-text search are ignored — the overview is an aggregate. +func (s *SqlStore) applyAgentNetworkUsageFilters(query *gorm.DB, filter agentNetworkTypes.AgentNetworkAccessLogFilter) *gorm.DB { + if filter.UserID != nil { + query = query.Where("user_id = ?", *filter.UserID) + } + if filter.SessionID != nil { + query = query.Where("session_id = ?", *filter.SessionID) + } + if len(filter.ProviderIDs) > 0 { + query = query.Where("resolved_provider_id IN ?", filter.ProviderIDs) + } + if len(filter.Models) > 0 { + query = query.Where("model IN ?", filter.Models) + } + if len(filter.GroupIDs) > 0 { + query = query.Where( + "id IN (SELECT usage_id FROM agent_network_request_usage_group WHERE group_id IN ?)", + filter.GroupIDs, + ) + } + if filter.StartDate != nil { + query = query.Where("timestamp >= ?", *filter.StartDate) + } + if filter.EndDate != nil { + query = query.Where("timestamp <= ?", *filter.EndDate) + } + return query +} + +// GetAgentNetworkAccessLogs retrieves flattened agent-network access logs for +// an account with server-side pagination, filtering and sorting. Authorising +// group ids are hydrated from the group child table for the returned page. +func (s *SqlStore) GetAgentNetworkAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLog, int64, error) { + var logs []*agentNetworkTypes.AgentNetworkAccessLog + var totalCount int64 + + countQuery := s.applyAgentNetworkAccessLogFilters( + s.db.Model(&agentNetworkTypes.AgentNetworkAccessLog{}).Where(accountIDCondition, accountID), + filter, + ) + if err := countQuery.Count(&totalCount).Error; err != nil { + log.WithContext(ctx).Errorf("failed to count agent-network access logs: %v", err) + return nil, 0, status.Errorf(status.Internal, "failed to count agent-network access logs") + } + + query := s.applyAgentNetworkAccessLogFilters( + s.db.Where(accountIDCondition, accountID), + filter, + ). + Order(filter.GetSortColumn() + " " + filter.GetSortOrder()). + Limit(filter.GetLimit()). + Offset(filter.GetOffset()) + + if lockStrength != LockingStrengthNone { + query = query.Clauses(clause.Locking{Strength: string(lockStrength)}) + } + + if err := query.Find(&logs).Error; err != nil { + log.WithContext(ctx).Errorf("failed to get agent-network access logs from store: %v", err) + return nil, 0, status.Errorf(status.Internal, "failed to get agent-network access logs from store") + } + + if err := s.hydrateAgentNetworkAccessLogGroups(ctx, accountID, logs); err != nil { + return nil, 0, err + } + + return logs, totalCount, nil +} + +// applyAgentNetworkAccessLogFilters applies the filter conditions to a query. +func (s *SqlStore) applyAgentNetworkAccessLogFilters(query *gorm.DB, filter agentNetworkTypes.AgentNetworkAccessLogFilter) *gorm.DB { + if filter.Search != nil { + p := "%" + *filter.Search + "%" + query = query.Where( + "id LIKE ? OR host LIKE ? OR path LIKE ? OR model LIKE ? OR user_id IN (SELECT id FROM users WHERE email LIKE ? OR name LIKE ?)", + p, p, p, p, p, p, + ) + } + if filter.UserID != nil { + query = query.Where("user_id = ?", *filter.UserID) + } + if filter.SessionID != nil { + query = query.Where("session_id = ?", *filter.SessionID) + } + if filter.Decision != nil { + query = query.Where("decision = ?", *filter.Decision) + } + if filter.PathPrefix != nil { + query = query.Where("path LIKE ?", *filter.PathPrefix+"%") + } + if len(filter.ProviderIDs) > 0 { + query = query.Where("resolved_provider_id IN ?", filter.ProviderIDs) + } + if len(filter.Models) > 0 { + query = query.Where("model IN ?", filter.Models) + } + if len(filter.GroupIDs) > 0 { + query = query.Where( + "id IN (SELECT log_id FROM agent_network_access_log_group WHERE group_id IN ?)", + filter.GroupIDs, + ) + } + if filter.StartDate != nil { + query = query.Where("timestamp >= ?", *filter.StartDate) + } + if filter.EndDate != nil { + query = query.Where("timestamp <= ?", *filter.EndDate) + } + return query +} + +// hydrateAgentNetworkAccessLogGroups loads the authorising group ids for the +// given page of entries and assigns them onto each entry's GroupIDs field. +func (s *SqlStore) hydrateAgentNetworkAccessLogGroups(ctx context.Context, accountID string, logs []*agentNetworkTypes.AgentNetworkAccessLog) error { + if len(logs) == 0 { + return nil + } + + ids := make([]string, 0, len(logs)) + for _, l := range logs { + ids = append(ids, l.ID) + } + + var rows []agentNetworkTypes.AgentNetworkAccessLogGroup + if err := s.db. + Where(accountIDCondition, accountID). + Where("log_id IN ?", ids). + Find(&rows).Error; err != nil { + log.WithContext(ctx).Errorf("failed to hydrate agent-network access log groups: %v", err) + return status.Errorf(status.Internal, "failed to hydrate agent-network access log groups") + } + + byLog := make(map[string][]string, len(logs)) + for _, r := range rows { + byLog[r.LogID] = append(byLog[r.LogID], r.GroupID) + } + for _, l := range logs { + l.GroupIDs = byLog[l.ID] + } + return nil +} + // GetAccountAccessLogs retrieves access logs for a given account with pagination and filtering func (s *SqlStore) GetAccountAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter accesslogs.AccessLogFilter) ([]*accesslogs.AccessLogEntry, int64, error) { var logs []*accesslogs.AccessLogEntry diff --git a/management/server/store/sql_store_agentnetwork.go b/management/server/store/sql_store_agentnetwork.go new file mode 100644 index 000000000..f0f1f8327 --- /dev/null +++ b/management/server/store/sql_store_agentnetwork.go @@ -0,0 +1,623 @@ +package store + +import ( + "context" + "errors" + "math" + "time" + + log "github.com/sirupsen/logrus" + "gorm.io/gorm" + "gorm.io/gorm/clause" + + agentNetworkTypes "github.com/netbirdio/netbird/management/server/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 +} + +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 +} + +// GetAgentNetworkSettingsByCluster returns every Settings row pinned to +// the given proxy cluster. Used by the bootstrap label generator to +// build the set of subdomains already taken on a cluster. +func (s *SqlStore) GetAgentNetworkSettingsByCluster(ctx context.Context, lockStrength LockingStrength, cluster 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, "cluster = ?", cluster) + if result.Error != nil { + log.WithContext(ctx).Errorf("failed to get agent network settings by cluster from store: %v", result.Error) + return nil, status.Errorf(status.Internal, "failed to get agent network settings by cluster from store") + } + + return settings, 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 +} + +// 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 +} diff --git a/management/server/store/sql_store_agentnetwork_accesslog_test.go b/management/server/store/sql_store_agentnetwork_accesslog_test.go new file mode 100644 index 000000000..c624f10a4 --- /dev/null +++ b/management/server/store/sql_store_agentnetwork_accesslog_test.go @@ -0,0 +1,202 @@ +package store + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + agentNetworkTypes "github.com/netbirdio/netbird/management/server/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, + CostUSD: 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, + CostUSD: 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, CostUSD: 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].CostUSD, 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) +} + +// 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") +} diff --git a/management/server/store/sql_store_agentnetwork_budgetrule_test.go b/management/server/store/sql_store_agentnetwork_budgetrule_test.go new file mode 100644 index 000000000..4e896e859 --- /dev/null +++ b/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/server/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, + Cluster: "eu.proxy.netbird.io", + Subdomain: "violet", + })) + + 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") +} diff --git a/management/server/store/store.go b/management/server/store/store.go index 066ab285d..70d201bb8 100644 --- a/management/server/store/store.go +++ b/management/server/store/store.go @@ -37,6 +37,7 @@ import ( "github.com/netbirdio/netbird/util" "github.com/netbirdio/netbird/util/crypt" + agentNetworkTypes "github.com/netbirdio/netbird/management/server/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" @@ -300,6 +301,11 @@ 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) + 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 @@ -329,6 +335,34 @@ type Store interface { GetProxyMetrics(ctx context.Context) (ProxyMetrics, 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) + GetAgentNetworkSettingsByCluster(ctx context.Context, lockStrength LockingStrength, cluster string) ([]*agentNetworkTypes.Settings, error) + SaveAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) 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 diff --git a/management/server/store/store_mock_agentnetwork.go b/management/server/store/store_mock_agentnetwork.go new file mode 100644 index 000000000..22e793a65 --- /dev/null +++ b/management/server/store/store_mock_agentnetwork.go @@ -0,0 +1,464 @@ +package store + +import ( + context "context" + reflect "reflect" + time "time" + + gomock "github.com/golang/mock/gomock" + + agentNetworkTypes "github.com/netbirdio/netbird/management/server/agentnetwork/types" +) + +// GetAllAgentNetworkProviders mocks base method. +func (m *MockStore) GetAllAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Provider, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAllAgentNetworkProviders", ctx, lockStrength) + ret0, _ := ret[0].([]*agentNetworkTypes.Provider) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAllAgentNetworkProviders indicates an expected call of GetAllAgentNetworkProviders. +func (mr *MockStoreMockRecorder) GetAllAgentNetworkProviders(ctx, lockStrength interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAllAgentNetworkProviders", reflect.TypeOf((*MockStore)(nil).GetAllAgentNetworkProviders), ctx, lockStrength) +} + +// GetAccountAgentNetworkProviders mocks base method. +func (m *MockStore) GetAccountAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Provider, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAccountAgentNetworkProviders", ctx, lockStrength, accountID) + ret0, _ := ret[0].([]*agentNetworkTypes.Provider) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAccountAgentNetworkProviders indicates an expected call of GetAccountAgentNetworkProviders. +func (mr *MockStoreMockRecorder) GetAccountAgentNetworkProviders(ctx, lockStrength, accountID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountAgentNetworkProviders", reflect.TypeOf((*MockStore)(nil).GetAccountAgentNetworkProviders), ctx, lockStrength, accountID) +} + +// GetAgentNetworkProviderByID mocks base method. +func (m *MockStore) GetAgentNetworkProviderByID(ctx context.Context, lockStrength LockingStrength, accountID, providerID string) (*agentNetworkTypes.Provider, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAgentNetworkProviderByID", ctx, lockStrength, accountID, providerID) + ret0, _ := ret[0].(*agentNetworkTypes.Provider) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAgentNetworkProviderByID indicates an expected call of GetAgentNetworkProviderByID. +func (mr *MockStoreMockRecorder) GetAgentNetworkProviderByID(ctx, lockStrength, accountID, providerID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkProviderByID", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkProviderByID), ctx, lockStrength, accountID, providerID) +} + +// SaveAgentNetworkProvider mocks base method. +func (m *MockStore) SaveAgentNetworkProvider(ctx context.Context, provider *agentNetworkTypes.Provider) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SaveAgentNetworkProvider", ctx, provider) + ret0, _ := ret[0].(error) + return ret0 +} + +// SaveAgentNetworkProvider indicates an expected call of SaveAgentNetworkProvider. +func (mr *MockStoreMockRecorder) SaveAgentNetworkProvider(ctx, provider interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkProvider", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkProvider), ctx, provider) +} + +// DeleteAgentNetworkProvider mocks base method. +func (m *MockStore) DeleteAgentNetworkProvider(ctx context.Context, accountID, providerID string) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteAgentNetworkProvider", ctx, accountID, providerID) + ret0, _ := ret[0].(error) + return ret0 +} + +// DeleteAgentNetworkProvider indicates an expected call of DeleteAgentNetworkProvider. +func (mr *MockStoreMockRecorder) DeleteAgentNetworkProvider(ctx, accountID, providerID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkProvider", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkProvider), ctx, accountID, providerID) +} + +// GetAccountAgentNetworkPolicies mocks base method. +func (m *MockStore) GetAccountAgentNetworkPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Policy, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAccountAgentNetworkPolicies", ctx, lockStrength, accountID) + ret0, _ := ret[0].([]*agentNetworkTypes.Policy) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAccountAgentNetworkPolicies indicates an expected call of GetAccountAgentNetworkPolicies. +func (mr *MockStoreMockRecorder) GetAccountAgentNetworkPolicies(ctx, lockStrength, accountID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountAgentNetworkPolicies", reflect.TypeOf((*MockStore)(nil).GetAccountAgentNetworkPolicies), ctx, lockStrength, accountID) +} + +// GetAgentNetworkPolicyByID mocks base method. +func (m *MockStore) GetAgentNetworkPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*agentNetworkTypes.Policy, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAgentNetworkPolicyByID", ctx, lockStrength, accountID, policyID) + ret0, _ := ret[0].(*agentNetworkTypes.Policy) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAgentNetworkPolicyByID indicates an expected call of GetAgentNetworkPolicyByID. +func (mr *MockStoreMockRecorder) GetAgentNetworkPolicyByID(ctx, lockStrength, accountID, policyID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkPolicyByID", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkPolicyByID), ctx, lockStrength, accountID, policyID) +} + +// SaveAgentNetworkPolicy mocks base method. +func (m *MockStore) SaveAgentNetworkPolicy(ctx context.Context, policy *agentNetworkTypes.Policy) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SaveAgentNetworkPolicy", ctx, policy) + ret0, _ := ret[0].(error) + return ret0 +} + +// SaveAgentNetworkPolicy indicates an expected call of SaveAgentNetworkPolicy. +func (mr *MockStoreMockRecorder) SaveAgentNetworkPolicy(ctx, policy interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkPolicy", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkPolicy), ctx, policy) +} + +// DeleteAgentNetworkPolicy mocks base method. +func (m *MockStore) DeleteAgentNetworkPolicy(ctx context.Context, accountID, policyID string) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteAgentNetworkPolicy", ctx, accountID, policyID) + ret0, _ := ret[0].(error) + return ret0 +} + +// DeleteAgentNetworkPolicy indicates an expected call of DeleteAgentNetworkPolicy. +func (mr *MockStoreMockRecorder) DeleteAgentNetworkPolicy(ctx, accountID, policyID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkPolicy", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkPolicy), ctx, accountID, policyID) +} + +// GetAccountAgentNetworkGuardrails mocks base method. +func (m *MockStore) GetAccountAgentNetworkGuardrails(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Guardrail, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAccountAgentNetworkGuardrails", ctx, lockStrength, accountID) + ret0, _ := ret[0].([]*agentNetworkTypes.Guardrail) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAccountAgentNetworkGuardrails indicates an expected call of GetAccountAgentNetworkGuardrails. +func (mr *MockStoreMockRecorder) GetAccountAgentNetworkGuardrails(ctx, lockStrength, accountID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountAgentNetworkGuardrails", reflect.TypeOf((*MockStore)(nil).GetAccountAgentNetworkGuardrails), ctx, lockStrength, accountID) +} + +// GetAgentNetworkGuardrailByID mocks base method. +func (m *MockStore) GetAgentNetworkGuardrailByID(ctx context.Context, lockStrength LockingStrength, accountID, guardrailID string) (*agentNetworkTypes.Guardrail, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAgentNetworkGuardrailByID", ctx, lockStrength, accountID, guardrailID) + ret0, _ := ret[0].(*agentNetworkTypes.Guardrail) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAgentNetworkGuardrailByID indicates an expected call of GetAgentNetworkGuardrailByID. +func (mr *MockStoreMockRecorder) GetAgentNetworkGuardrailByID(ctx, lockStrength, accountID, guardrailID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkGuardrailByID", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkGuardrailByID), ctx, lockStrength, accountID, guardrailID) +} + +// SaveAgentNetworkGuardrail mocks base method. +func (m *MockStore) SaveAgentNetworkGuardrail(ctx context.Context, guardrail *agentNetworkTypes.Guardrail) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SaveAgentNetworkGuardrail", ctx, guardrail) + ret0, _ := ret[0].(error) + return ret0 +} + +// SaveAgentNetworkGuardrail indicates an expected call of SaveAgentNetworkGuardrail. +func (mr *MockStoreMockRecorder) SaveAgentNetworkGuardrail(ctx, guardrail interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkGuardrail", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkGuardrail), ctx, guardrail) +} + +// DeleteAgentNetworkGuardrail mocks base method. +func (m *MockStore) DeleteAgentNetworkGuardrail(ctx context.Context, accountID, guardrailID string) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteAgentNetworkGuardrail", ctx, accountID, guardrailID) + ret0, _ := ret[0].(error) + return ret0 +} + +// DeleteAgentNetworkGuardrail indicates an expected call of DeleteAgentNetworkGuardrail. +func (mr *MockStoreMockRecorder) DeleteAgentNetworkGuardrail(ctx, accountID, guardrailID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkGuardrail", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkGuardrail), ctx, accountID, guardrailID) +} + +// GetAgentNetworkSettings mocks base method. +func (m *MockStore) GetAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*agentNetworkTypes.Settings, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAgentNetworkSettings", ctx, lockStrength, accountID) + ret0, _ := ret[0].(*agentNetworkTypes.Settings) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAgentNetworkSettings indicates an expected call of GetAgentNetworkSettings. +func (mr *MockStoreMockRecorder) GetAgentNetworkSettings(ctx, lockStrength, accountID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkSettings", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkSettings), ctx, lockStrength, accountID) +} + +// GetAgentNetworkSettingsByCluster mocks base method. +func (m *MockStore) GetAgentNetworkSettingsByCluster(ctx context.Context, lockStrength LockingStrength, cluster string) ([]*agentNetworkTypes.Settings, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAgentNetworkSettingsByCluster", ctx, lockStrength, cluster) + ret0, _ := ret[0].([]*agentNetworkTypes.Settings) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAgentNetworkSettingsByCluster indicates an expected call of GetAgentNetworkSettingsByCluster. +func (mr *MockStoreMockRecorder) GetAgentNetworkSettingsByCluster(ctx, lockStrength, cluster interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkSettingsByCluster", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkSettingsByCluster), ctx, lockStrength, cluster) +} + +// SaveAgentNetworkSettings mocks base method. +func (m *MockStore) SaveAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SaveAgentNetworkSettings", ctx, settings) + ret0, _ := ret[0].(error) + return ret0 +} + +// SaveAgentNetworkSettings indicates an expected call of SaveAgentNetworkSettings. +func (mr *MockStoreMockRecorder) SaveAgentNetworkSettings(ctx, settings interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkSettings", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkSettings), ctx, settings) +} + +// IncrementAgentNetworkConsumption mocks base method. +func (m *MockStore) IncrementAgentNetworkConsumption(ctx context.Context, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time, tokensIn, tokensOut int64, costUSD float64) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "IncrementAgentNetworkConsumption", ctx, accountID, kind, dimID, windowSeconds, windowStart, tokensIn, tokensOut, costUSD) + ret0, _ := ret[0].(error) + return ret0 +} + +// IncrementAgentNetworkConsumption indicates an expected call of IncrementAgentNetworkConsumption. +func (mr *MockStoreMockRecorder) IncrementAgentNetworkConsumption(ctx, accountID, kind, dimID, windowSeconds, windowStart, tokensIn, tokensOut, costUSD interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IncrementAgentNetworkConsumption", reflect.TypeOf((*MockStore)(nil).IncrementAgentNetworkConsumption), ctx, accountID, kind, dimID, windowSeconds, windowStart, tokensIn, tokensOut, costUSD) +} + +// GetAgentNetworkConsumption mocks base method. +func (m *MockStore) GetAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time) (*agentNetworkTypes.Consumption, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAgentNetworkConsumption", ctx, lockStrength, accountID, kind, dimID, windowSeconds, windowStart) + ret0, _ := ret[0].(*agentNetworkTypes.Consumption) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAgentNetworkConsumption indicates an expected call of GetAgentNetworkConsumption. +func (mr *MockStoreMockRecorder) GetAgentNetworkConsumption(ctx, lockStrength, accountID, kind, dimID, windowSeconds, windowStart interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkConsumption", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkConsumption), ctx, lockStrength, accountID, kind, dimID, windowSeconds, windowStart) +} + +// GetAgentNetworkConsumptionBatch mocks base method. +func (m *MockStore) GetAgentNetworkConsumptionBatch(ctx context.Context, lockStrength LockingStrength, accountID string, keys []agentNetworkTypes.ConsumptionKey) (map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAgentNetworkConsumptionBatch", ctx, lockStrength, accountID, keys) + ret0, _ := ret[0].(map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAgentNetworkConsumptionBatch indicates an expected call of GetAgentNetworkConsumptionBatch. +func (mr *MockStoreMockRecorder) GetAgentNetworkConsumptionBatch(ctx, lockStrength, accountID, keys interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkConsumptionBatch", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkConsumptionBatch), ctx, lockStrength, accountID, keys) +} + +// IncrementAgentNetworkConsumptionBatch mocks base method. +func (m *MockStore) IncrementAgentNetworkConsumptionBatch(ctx context.Context, accountID string, keys []agentNetworkTypes.ConsumptionKey, tokensIn, tokensOut int64, costUSD float64) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "IncrementAgentNetworkConsumptionBatch", ctx, accountID, keys, tokensIn, tokensOut, costUSD) + ret0, _ := ret[0].(error) + return ret0 +} + +// IncrementAgentNetworkConsumptionBatch indicates an expected call of IncrementAgentNetworkConsumptionBatch. +func (mr *MockStoreMockRecorder) IncrementAgentNetworkConsumptionBatch(ctx, accountID, keys, tokensIn, tokensOut, costUSD interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IncrementAgentNetworkConsumptionBatch", reflect.TypeOf((*MockStore)(nil).IncrementAgentNetworkConsumptionBatch), ctx, accountID, keys, tokensIn, tokensOut, costUSD) +} + +// ListAgentNetworkConsumption mocks base method. +func (m *MockStore) ListAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Consumption, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListAgentNetworkConsumption", ctx, lockStrength, accountID) + ret0, _ := ret[0].([]*agentNetworkTypes.Consumption) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListAgentNetworkConsumption indicates an expected call of ListAgentNetworkConsumption. +func (mr *MockStoreMockRecorder) ListAgentNetworkConsumption(ctx, lockStrength, accountID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAgentNetworkConsumption", reflect.TypeOf((*MockStore)(nil).ListAgentNetworkConsumption), ctx, lockStrength, accountID) +} + +// GetAccountAgentNetworkBudgetRules mocks base method. +func (m *MockStore) GetAccountAgentNetworkBudgetRules(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.AccountBudgetRule, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAccountAgentNetworkBudgetRules", ctx, lockStrength, accountID) + ret0, _ := ret[0].([]*agentNetworkTypes.AccountBudgetRule) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAccountAgentNetworkBudgetRules indicates an expected call of GetAccountAgentNetworkBudgetRules. +func (mr *MockStoreMockRecorder) GetAccountAgentNetworkBudgetRules(ctx, lockStrength, accountID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountAgentNetworkBudgetRules", reflect.TypeOf((*MockStore)(nil).GetAccountAgentNetworkBudgetRules), ctx, lockStrength, accountID) +} + +// GetAgentNetworkBudgetRuleByID mocks base method. +func (m *MockStore) GetAgentNetworkBudgetRuleByID(ctx context.Context, lockStrength LockingStrength, accountID, ruleID string) (*agentNetworkTypes.AccountBudgetRule, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAgentNetworkBudgetRuleByID", ctx, lockStrength, accountID, ruleID) + ret0, _ := ret[0].(*agentNetworkTypes.AccountBudgetRule) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAgentNetworkBudgetRuleByID indicates an expected call of GetAgentNetworkBudgetRuleByID. +func (mr *MockStoreMockRecorder) GetAgentNetworkBudgetRuleByID(ctx, lockStrength, accountID, ruleID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkBudgetRuleByID", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkBudgetRuleByID), ctx, lockStrength, accountID, ruleID) +} + +// SaveAgentNetworkBudgetRule mocks base method. +func (m *MockStore) SaveAgentNetworkBudgetRule(ctx context.Context, rule *agentNetworkTypes.AccountBudgetRule) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SaveAgentNetworkBudgetRule", ctx, rule) + ret0, _ := ret[0].(error) + return ret0 +} + +// SaveAgentNetworkBudgetRule indicates an expected call of SaveAgentNetworkBudgetRule. +func (mr *MockStoreMockRecorder) SaveAgentNetworkBudgetRule(ctx, rule interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkBudgetRule", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkBudgetRule), ctx, rule) +} + +// DeleteAgentNetworkBudgetRule mocks base method. +func (m *MockStore) DeleteAgentNetworkBudgetRule(ctx context.Context, accountID, ruleID string) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteAgentNetworkBudgetRule", ctx, accountID, ruleID) + ret0, _ := ret[0].(error) + return ret0 +} + +// DeleteAgentNetworkBudgetRule indicates an expected call of DeleteAgentNetworkBudgetRule. +func (mr *MockStoreMockRecorder) DeleteAgentNetworkBudgetRule(ctx, accountID, ruleID interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkBudgetRule", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkBudgetRule), ctx, accountID, ruleID) +} + +// CreateAgentNetworkAccessLog mocks base method. +func (m *MockStore) CreateAgentNetworkAccessLog(ctx context.Context, entry *agentNetworkTypes.AgentNetworkAccessLog, groups []agentNetworkTypes.AgentNetworkAccessLogGroup) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CreateAgentNetworkAccessLog", ctx, entry, groups) + ret0, _ := ret[0].(error) + return ret0 +} + +// CreateAgentNetworkAccessLog indicates an expected call of CreateAgentNetworkAccessLog. +func (mr *MockStoreMockRecorder) CreateAgentNetworkAccessLog(ctx, entry, groups interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateAgentNetworkAccessLog", reflect.TypeOf((*MockStore)(nil).CreateAgentNetworkAccessLog), ctx, entry, groups) +} + +// CreateAgentNetworkUsage mocks base method. +func (m *MockStore) CreateAgentNetworkUsage(ctx context.Context, usage *agentNetworkTypes.AgentNetworkUsage, groups []agentNetworkTypes.AgentNetworkUsageGroup) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CreateAgentNetworkUsage", ctx, usage, groups) + ret0, _ := ret[0].(error) + return ret0 +} + +// CreateAgentNetworkUsage indicates an expected call of CreateAgentNetworkUsage. +func (mr *MockStoreMockRecorder) CreateAgentNetworkUsage(ctx, usage, groups interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateAgentNetworkUsage", reflect.TypeOf((*MockStore)(nil).CreateAgentNetworkUsage), ctx, usage, groups) +} + +// GetAgentNetworkAccessLogs mocks base method. +func (m *MockStore) GetAgentNetworkAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLog, int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAgentNetworkAccessLogs", ctx, lockStrength, accountID, filter) + ret0, _ := ret[0].([]*agentNetworkTypes.AgentNetworkAccessLog) + ret1, _ := ret[1].(int64) + ret2, _ := ret[2].(error) + return ret0, ret1, ret2 +} + +// GetAgentNetworkAccessLogs indicates an expected call of GetAgentNetworkAccessLogs. +func (mr *MockStoreMockRecorder) GetAgentNetworkAccessLogs(ctx, lockStrength, accountID, filter interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkAccessLogs", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkAccessLogs), ctx, lockStrength, accountID, filter) +} + +// GetAgentNetworkUsageRows mocks base method. +func (m *MockStore) GetAgentNetworkUsageRows(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkUsage, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAgentNetworkUsageRows", ctx, lockStrength, accountID, filter) + ret0, _ := ret[0].([]*agentNetworkTypes.AgentNetworkUsage) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAgentNetworkUsageRows indicates an expected call of GetAgentNetworkUsageRows. +func (mr *MockStoreMockRecorder) GetAgentNetworkUsageRows(ctx, lockStrength, accountID, filter interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkUsageRows", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkUsageRows), ctx, lockStrength, accountID, filter) +} + +// DeleteOldAgentNetworkAccessLogs mocks base method. +func (m *MockStore) DeleteOldAgentNetworkAccessLogs(ctx context.Context, accountID string, olderThan time.Time) (int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteOldAgentNetworkAccessLogs", ctx, accountID, olderThan) + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// DeleteOldAgentNetworkAccessLogs indicates an expected call of DeleteOldAgentNetworkAccessLogs. +func (mr *MockStoreMockRecorder) DeleteOldAgentNetworkAccessLogs(ctx, accountID, olderThan interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOldAgentNetworkAccessLogs", reflect.TypeOf((*MockStore)(nil).DeleteOldAgentNetworkAccessLogs), ctx, accountID, olderThan) +} + +// GetAllAgentNetworkSettings mocks base method. +func (m *MockStore) GetAllAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Settings, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAllAgentNetworkSettings", ctx, lockStrength) + ret0, _ := ret[0].([]*agentNetworkTypes.Settings) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAllAgentNetworkSettings indicates an expected call of GetAllAgentNetworkSettings. +func (mr *MockStoreMockRecorder) GetAllAgentNetworkSettings(ctx, lockStrength interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAllAgentNetworkSettings", reflect.TypeOf((*MockStore)(nil).GetAllAgentNetworkSettings), ctx, lockStrength) +}