package agentnetwork import ( "context" "errors" "fmt" "math/rand" "slices" "strings" "sync" "time" log "github.com/sirupsen/logrus" "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/labelgen" "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types" "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/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) ListAccessLogSessions(ctx context.Context, accountID, userID string, filter types.AgentNetworkAccessLogFilter) ([]*types.AgentNetworkAccessLogSession, 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 and the model it extracted. type PolicySelectionInput struct { AccountID string UserID string GroupIDs []string ProviderID string // Model is the already-normalised upstream model id the proxy extracted // (parser strips Bedrock region/version, Vertex @version), so a // case-insensitive compare suffices. Empty = undetermined → not permitted // (fail closed). Model string } // PolicySelectionResult names the policy that "pays" for this request // 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, modules.AgentNetworkProviders, 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, modules.AgentNetworkProviders, 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, modules.AgentNetworkProviders, operations.Create); err != nil { return nil, err } if strings.TrimSpace(bootstrapCluster) != "" { if err := m.requireSettingsBootstrapPermission(ctx, provider.AccountID, userID); 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, m.store, 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, modules.AgentNetworkProviders, 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, modules.AgentNetworkProviders, 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, modules.AgentNetworkPolicies, 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, modules.AgentNetworkPolicies, 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, modules.AgentNetworkPolicies, 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, modules.AgentNetworkPolicies, 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, modules.AgentNetworkPolicies, 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, modules.AgentNetworkGuardrails, 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, modules.AgentNetworkGuardrails, 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, modules.AgentNetworkGuardrails, 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, modules.AgentNetworkGuardrails, 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, modules.AgentNetworkGuardrails, 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, modules.AgentNetworkBudgets, 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, modules.AgentNetworkBudgets, 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, modules.AgentNetworkBudgets, 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, modules.AgentNetworkBudgets, 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, modules.AgentNetworkBudgets, 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 replaces the mutable account-level settings — the collection // toggles and retention — on the account's row. When the account has no // settings row yet, a non-empty settings.Cluster bootstraps one (same path as // first provider create); without it the update fails with NotFound. On an // existing row the cluster and subdomain are immutable: a differing // settings.Cluster is rejected rather than silently ignored so callers never // observe a value other than what they sent. 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, modules.AgentNetworkSettings, operations.Update); err != nil { return nil, err } requestedCluster := strings.TrimSpace(settings.Cluster) // The row lock from LockingStrengthUpdate only holds for the duration of // the surrounding transaction, so the read, the cluster-immutability // check, and the save must share one — otherwise concurrent PUTs could // interleave between them. var updated *types.Settings err := m.store.ExecuteInTransaction(ctx, func(tx store.Store) error { existing, err := tx.GetAgentNetworkSettings(ctx, store.LockingStrengthUpdate, settings.AccountID) switch { case err == nil: if requestedCluster != "" && requestedCluster != existing.Cluster { return status.Errorf(status.InvalidArgument, "cluster is immutable once assigned (current: %s)", existing.Cluster) } case isNotFound(err): if requestedCluster == "" { return status.Errorf(status.NotFound, "agent network settings have not been bootstrapped yet; pass cluster to bootstrap them, or create a provider with bootstrap_cluster set") } // Bootstrapping pins the cluster and subdomain — a settings // create on top of the update the caller already passed, matching // the gate on the provider-create bootstrap path. if err := m.requirePermission(ctx, settings.AccountID, userID, modules.AgentNetworkSettings, operations.Create); err != nil { return err } existing, err = m.bootstrapSettingsIfNeeded(ctx, tx, settings.AccountID, requestedCluster) if err != nil { return err } default: return 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 := tx.SaveAgentNetworkSettings(ctx, existing); err != nil { return fmt.Errorf("save agent network settings: %w", err) } updated = existing return nil }) if err != nil { return nil, err } m.accountManager.StoreEvent(ctx, userID, settings.AccountID, settings.AccountID, activity.AgentNetworkSettingsUpdated, map[string]any{ "log_collection": updated.EnableLogCollection, "prompt_collection": updated.EnablePromptCollection, "redact_pii": updated.RedactPii, }) m.reconcile(ctx, settings.AccountID) return updated, nil } // isNotFound reports whether err is a status.NotFound error. func isNotFound(err error) bool { var sErr *status.Error return errors.As(err, &sErr) && sErr.Type() == status.NotFound } // 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. When no // row has been bootstrapped yet, the defaults are returned (without // persisting) with cluster and subdomain empty — settings always read as an // object, like the account and DNS settings endpoints. func (m *managerImpl) GetSettings(ctx context.Context, accountID, userID string) (*types.Settings, error) { if err := m.requirePermission(ctx, accountID, userID, modules.AgentNetworkSettings, operations.Read); err != nil { return nil, err } settings, err := m.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID) switch { case err == nil: return settings, nil case isNotFound(err): return types.DefaultSettings(accountID), nil default: return nil, err } } // requireSettingsBootstrapPermission gates the one-time settings bootstrap a // first provider create performs. Pinning the account's cluster and subdomain // is a settings write, so it needs the settings permission on top of the // provider one. No-op once the settings row exists. func (m *managerImpl) requireSettingsBootstrapPermission(ctx context.Context, accountID, userID string) error { _, err := m.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID) if err == nil { return nil } if !isNotFound(err) { return fmt.Errorf("get agent network settings: %w", err) } return m.requirePermission(ctx, accountID, userID, modules.AgentNetworkSettings, operations.Create) } // 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. st is the store to // operate on — pass the transaction store when calling from within one. func (m *managerImpl) bootstrapSettingsIfNeeded(ctx context.Context, st store.Store, 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 := st.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID) if err == nil { return existing, nil } if !isNotFound(err) { return nil, fmt.Errorf("get agent network settings: %w", err) } siblings, err := st.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.DefaultSettings(accountID) settings.Cluster = providerCluster settings.Subdomain = subdomain settings.CreatedAt = now settings.UpdatedAt = now if err := st.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, modules.AgentNetworkUsage, 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, modules.AgentNetworkLogs, operations.Read); err != nil { return nil, 0, err } return m.store.GetAgentNetworkAccessLogs(ctx, store.LockingStrengthNone, accountID, filter) } // ListAccessLogSessions returns a paginated, server-side-filtered page of // agent-network access logs grouped by session, plus the total number of // sessions matching the filter. func (m *managerImpl) ListAccessLogSessions(ctx context.Context, accountID, userID string, filter types.AgentNetworkAccessLogFilter) ([]*types.AgentNetworkAccessLogSession, int64, error) { if err := m.requirePermission(ctx, accountID, userID, modules.AgentNetworkLogs, operations.Read); err != nil { return nil, 0, err } return m.store.GetAgentNetworkAccessLogSessions(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, modules.AgentNetworkUsage, 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, module modules.Module, op operations.Operation) error { ok, _, err := m.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, module, 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, accountID, _ string) (*types.Settings, error) { return types.DefaultSettings(accountID), nil } 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) ListAccessLogSessions(_ context.Context, _, _ string, _ types.AgentNetworkAccessLogFilter) ([]*types.AgentNetworkAccessLogSession, 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 }