mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-19 13:19:06 +02:00
[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.
This commit is contained in:
mlsmaycon
@@ -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: "<resource>.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.<region>.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: "<region>-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::<base64 SA>").
|
||||
// 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
|
||||
}
|
||||
@@ -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
|
||||
// `-<fallbackSuffix>` 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)
|
||||
}
|
||||
@@ -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 <word>-<suffix>; 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")
|
||||
}
|
||||
@@ -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",
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 <id>" or "account rule <id>") 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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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 (`<slug>.<cluster>`) 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 ""
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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")
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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" }
|
||||
@@ -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
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 (`<subdomain>.<cluster>`).
|
||||
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:
|
||||
// `<subdomain>.<cluster>`.
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -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" }
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
Reference in New Issue
Block a user