package agentnetwork import ( "context" "errors" "fmt" "math/rand" "slices" "strings" "sync" "time" "unicode" "unicode/utf8" log "github.com/sirupsen/logrus" "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog" "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) DiscoverProviderModels(ctx context.Context, accountID, userID, providerID string) (*types.ModelDiscoveryResult, 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. 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 // discoveryAttempts provides lightweight per-provider admission control for // the explicit endpoint probe. It prevents duplicate clicks from occupying // multiple proxy control-stream requests at once and adds a short cooldown // after each attempt. discoveryMu sync.Mutex discoveryAttempts map[string]modelDiscoveryAttempt } type modelDiscoveryAttempt struct { inFlight bool lastStarted time.Time } const ( modelDiscoveryTimeout = 10 * time.Second modelDiscoveryCooldown = 2 * time.Second maxDiscoveredModels = 500 maxDiscoveredModelLen = 512 ) // 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())), discoveryAttempts: make(map[string]modelDiscoveryAttempt), } } 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) } // DiscoverProviderModels asks a capable proxy in the account's selected // cluster to query the persisted provider endpoint. The browser supplies only // the provider id: URL, TLS policy, and credential are loaded here so a caller // cannot turn this operation into an arbitrary network probe. func (m *managerImpl) DiscoverProviderModels(ctx context.Context, accountID, userID, providerID string) (*types.ModelDiscoveryResult, error) { if err := m.requirePermission(ctx, accountID, userID, operations.Update); err != nil { return nil, err } providerID = strings.TrimSpace(providerID) if providerID == "" { return nil, status.Errorf(status.InvalidArgument, "provider ID is required") } provider, err := m.store.GetAgentNetworkProviderByID(ctx, store.LockingStrengthNone, accountID, providerID) if err != nil { return nil, err } entry, ok := catalog.Lookup(provider.ProviderID) if !ok { return nil, status.Errorf(status.PreconditionFailed, "provider references an unknown catalog provider") } if entry.ModelDiscovery == nil { return nil, status.Errorf(status.PreconditionFailed, "provider type %q does not support model discovery", provider.ProviderID) } if m.proxyController == nil { return nil, status.Errorf(status.PreconditionFailed, "model discovery is unavailable") } settings, err := m.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID) if err != nil { var statusErr *status.Error if errors.As(err, &statusErr) && statusErr.Type() == status.NotFound { return nil, status.Errorf(status.PreconditionFailed, "configure an Agent Network proxy cluster before discovering models") } return nil, err } cluster := strings.TrimSpace(settings.Cluster) if cluster == "" { return nil, status.Errorf(status.PreconditionFailed, "configure an Agent Network proxy cluster before discovering models") } authHeaderName, authHeaderValue, gcpKey, err := providerAuthHeader(provider) if err != nil { return nil, status.Errorf(status.PreconditionFailed, "provider authentication is not configured correctly") } if gcpKey != "" { return nil, status.Errorf(status.PreconditionFailed, "provider authentication is not supported for model discovery") } attemptKey := accountID + "\x00" + provider.ID if !m.beginModelDiscovery(attemptKey, time.Now()) { return nil, status.Errorf(status.TooManyRequests, "model discovery is already running or was requested too recently") } defer m.finishModelDiscovery(attemptKey) probeCtx, cancel := context.WithTimeout(ctx, modelDiscoveryTimeout) defer cancel() probeResult, err := m.proxyController.DiscoverModels(probeCtx, accountID, cluster, &proto.ModelDiscoveryRequest{ UpstreamUrl: strings.TrimSpace(provider.UpstreamURL), AuthHeaderName: authHeaderName, AuthHeaderValue: authHeaderValue, SkipTlsVerify: provider.SkipTLSVerification, OllamaFallback: entry.ModelDiscovery.OllamaFallback, }) if err != nil { if errors.Is(err, context.DeadlineExceeded) || errors.Is(probeCtx.Err(), context.DeadlineExceeded) { return nil, status.Errorf(status.PreconditionFailed, "model discovery timed out") } if _, ok := status.FromError(err); ok { return nil, err } return nil, status.Errorf(status.PreconditionFailed, "model discovery could not be completed") } if probeResult == nil { return nil, status.Errorf(status.PreconditionFailed, "proxy returned an empty model discovery response") } if strings.TrimSpace(probeResult.Error) != "" { message := safeModelDiscoveryText(probeResult.Error, 256) if message == "" { message = "proxy reported a discovery failure" } requestID := safeModelDiscoveryText(probeResult.RequestId, 128) if requestID != "" { return nil, status.Errorf(status.PreconditionFailed, "model discovery failed: %s (request_id: %s)", message, requestID) } return nil, status.Errorf(status.PreconditionFailed, "model discovery failed: %s", message) } requestID := safeModelDiscoveryText(probeResult.RequestId, 128) if requestID == "" { return nil, status.Errorf(status.PreconditionFailed, "proxy returned an uncorrelated model discovery response") } source := strings.TrimSpace(probeResult.Source) switch source { case "openai_v1_models", "ollama_api_tags": default: return nil, status.Errorf(status.PreconditionFailed, "proxy returned an unsupported model discovery response") } return &types.ModelDiscoveryResult{ RequestID: requestID, Source: source, ProxyCluster: cluster, Models: normalizeDiscoveredModels(probeResult.Models), }, nil } func (m *managerImpl) beginModelDiscovery(key string, now time.Time) bool { m.discoveryMu.Lock() defer m.discoveryMu.Unlock() if m.discoveryAttempts == nil { m.discoveryAttempts = make(map[string]modelDiscoveryAttempt) } attempt := m.discoveryAttempts[key] if attempt.inFlight || (!attempt.lastStarted.IsZero() && now.Sub(attempt.lastStarted) < modelDiscoveryCooldown) { return false } attempt.inFlight = true attempt.lastStarted = now m.discoveryAttempts[key] = attempt return true } func (m *managerImpl) finishModelDiscovery(key string) { m.discoveryMu.Lock() attempt, ok := m.discoveryAttempts[key] if !ok { m.discoveryMu.Unlock() return } attempt.inFlight = false m.discoveryAttempts[key] = attempt m.discoveryMu.Unlock() remaining := time.Until(attempt.lastStarted.Add(modelDiscoveryCooldown)) if remaining <= 0 { m.expireModelDiscoveryAttempt(key, attempt.lastStarted) return } time.AfterFunc(remaining, func() { m.expireModelDiscoveryAttempt(key, attempt.lastStarted) }) } func (m *managerImpl) expireModelDiscoveryAttempt(key string, lastStarted time.Time) { m.discoveryMu.Lock() defer m.discoveryMu.Unlock() attempt, ok := m.discoveryAttempts[key] if ok && !attempt.inFlight && attempt.lastStarted.Equal(lastStarted) { delete(m.discoveryAttempts, key) } } func normalizeDiscoveredModels(models []*proto.ModelDiscoveryModel) []types.DiscoveredModel { out := make([]types.DiscoveredModel, 0, min(len(models), maxDiscoveredModels)) seen := make(map[string]struct{}, min(len(models), maxDiscoveredModels)) for _, model := range models { if model == nil { continue } id := safeModelDiscoveryText(model.Id, maxDiscoveredModelLen) if id == "" || len(id) > maxDiscoveredModelLen { continue } if _, ok := seen[id]; ok { continue } label := safeModelDiscoveryText(model.Label, maxDiscoveredModelLen) if label == "" || len(label) > maxDiscoveredModelLen { label = id } seen[id] = struct{}{} out = append(out, types.DiscoveredModel{ID: id, Label: label}) if len(out) == maxDiscoveredModels { break } } slices.SortFunc(out, func(a, b types.DiscoveredModel) int { return strings.Compare(a.ID, b.ID) }) return out } func safeModelDiscoveryText(value string, maxLen int) string { value = strings.TrimSpace(value) if value == "" || len(value) > maxLen || !utf8.ValidString(value) { return "" } for _, r := range value { if unicode.IsControl(r) { return "" } } return value } // 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 } if err := prepareProviderAPIKey(provider, nil); err != nil { return nil, err } 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) } if err := prepareProviderAPIKey(provider, existing); err != nil { return nil, err } // 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 } // prepareProviderAPIKey applies catalog-owned authentication semantics before // a provider is persisted. existing is nil on create. On update, omission // preserves a key only when the provider type is unchanged; switching types // never carries an old provider's secret into the new upstream. func prepareProviderAPIKey(provider, existing *types.Provider) error { entry, ok := catalog.Lookup(provider.ProviderID) if !ok { return status.Errorf(status.InvalidArgument, "provider_id %q is not a known catalog provider", provider.ProviderID) } return prepareProviderAPIKeyForEntry(provider, existing, entry) } func prepareProviderAPIKeyForEntry(provider, existing *types.Provider, entry catalog.Provider) error { keyProvided := provider.APIKeyProvided || provider.APIKey != "" providerChanged := existing != nil && provider.ProviderID != existing.ProviderID if existing != nil && !keyProvided && !providerChanged && entry.EffectiveAuthMode() != catalog.AuthModeNone { provider.APIKey = existing.APIKey } // Normalize after a preserved value is restored so a legacy // whitespace-only credential cannot pass manager validation and then fail // later during synthesis. if strings.TrimSpace(provider.APIKey) == "" { provider.APIKey = "" } switch entry.EffectiveAuthMode() { case catalog.AuthModeRequired: if provider.APIKey == "" { return status.Errorf(status.InvalidArgument, "api_key is required for provider_id %q", provider.ProviderID) } case catalog.AuthModeOptional: // Empty is a valid credential state. case catalog.AuthModeNone: if keyProvided && provider.APIKey != "" { return status.Errorf(status.InvalidArgument, "api_key is not supported for provider_id %q", provider.ProviderID) } // Clear any stale credential already stored for a provider whose // catalog semantics changed to no authentication. provider.APIKey = "" default: return status.Errorf(status.InvalidArgument, "provider_id %q has an invalid authentication mode", provider.ProviderID) } return 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) } // 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, 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, 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) DiscoverProviderModels(_ context.Context, _, _, _ string) (*types.ModelDiscoveryResult, error) { return nil, status.Errorf(status.PreconditionFailed, "model discovery is unavailable") } 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) 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 }