[management,proxy] Agent network: per-account LLM gateway (policy, metering, multi-provider) (#6555)

* [agent-network] Shared proto, OpenAPI schema, and generated types * [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. * [management] Fix agent-network proxy-peer fan-out on affected-peer recompute The affected-peers resolver loaded only persisted reverse-proxy services, but agent-network services are synthesized on demand and never persisted. As a result the embedded proxy peer was never folded into the affected set when a client's group changed, so the proxy received no network-map update for a newly authorised client and rejected its handshake until a full resync (restart). loadProxyServices now merges the synthesized agent-network services (injected via a registration hook to avoid an import cycle), so proxy peers learn newly authorised clients immediately. * [proxy] Reverse-proxy middleware framework, chain, and request plumbing The per-target middleware chain (slots, dispatcher, mutation gate, metadata merger), body capture, access-log terminal sink, and the proxy wiring that builds + runs chains for synthesized agent-network services. * [proxy] LLM parsers, pricing, and builtin middlewares (OpenAI, Anthropic, Vertex AI, AWS Bedrock) Request/response parsers and SSE/event-stream metering, the embedded pricing table, and the builtin middleware set: request parser, router, policy limit-check/record, cost meter, guardrail, identity inject, response parser. Includes the path-routed providers — Google Vertex AI (keyfile:: service-account OAuth minting) and AWS Bedrock (bearer auth, invoke/converse/streaming, optional /bedrock prefix) — plus the Models allowlist and unmeterable-publisher deny. * [proxy] IPv6 in-place apply and TCP accept-loop hardening on netstack listeners * [agent-network] End-to-end test suite, module docs, and deployment preset * [agent-network] Fix codespell typos and exclude false positives - labelgen word pool: vermillion -> vermilion, racoon -> raccoon. - codespell ignore list: add flate (Go compress/flate package), recordin (a test-local identifier), and unparseable (a valid alternative spelling used consistently across identifiers + a metadata-value constant). * [management] Set LastSeen on injected proxy peer in realstack test (MySQL strict-mode) The injected embedded proxy peer had a PeerStatus with a zero LastSeen, which serializes to '0000-00-00' and is rejected by MySQL in strict mode (SQLite tolerates it). Set LastSeen to a valid time so SaveAccount succeeds on both engines. * [agent-network] Remove e2e shell-script suite from this branch The end-to-end shell scripts under scripts/e2e/ are maintained in a separate testing suite and are not part of this change set. * [agent-network] Polish module docs: remove internal review scaffolding, fix links, verify diagrams Strip PR-review framing, commit references, absolute paths, and stale internal references from the agent-network module docs; fix broken relative links; verify all diagrams against the current architecture. Remove the internal AI-reviewer prompt file. * [management] Refine session expiration handling to support 3-state encoding for SSO deadlines * [agent-network] Relocate agentnetwork package to internals/modules Move management/server/agentnetwork (and its catalog/, labelgen/, types/ subpackages) to management/internals/modules/agentnetwork, alongside the reverse-proxy module, and rewrite all importers. Pure relocation: package names, the synthesizer + affectedpeers registration hook, and store access (shared store.Store) are unchanged, so no import cycle is introduced (affectedpeers still depends only on the agentnetwork/types leaf). * [agent-network] Co-locate HTTP handlers in the module (RegisterEndpoints) Move the agent-network HTTP handlers from server/http/handlers/agentnetwork into the module at internals/modules/agentnetwork/handlers (package handlers) and rename the entrypoint AddEndpoints -> RegisterEndpoints, matching the reverse-proxy module convention. Wiring in http/handler.go updated accordingly.
2026-10-01 19:19:07 +02:00 · 2026-06-27 13:41:00 +02:00
parent 615631567a
commit b416063bcc
187 changed files with 36835 additions and 660 deletions
@@ -2,12 +2,15 @@ package manager

 import (
 	"context"
+	"math"
+	"strconv"
 	"strings"
 	"time"

 	log "github.com/sirupsen/logrus"

 	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
+	agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
 	"github.com/netbirdio/netbird/management/server/geolocation"
 	"github.com/netbirdio/netbird/management/server/permissions"
 	"github.com/netbirdio/netbird/management/server/permissions/modules"
@@ -16,6 +19,28 @@ import (
 	"github.com/netbirdio/netbird/shared/management/status"
 )

+// Metadata keys the proxy stamps on agent-network access-log entries. These
+// mirror the constants in proxy/internal/middleware/keys.go and form the wire
+// contract between the proxy and management; management flattens them into
+// queryable columns. Keep in sync with the proxy side.
+const (
+	metaKeyProvider           = "llm.provider"
+	metaKeyModel              = "llm.model"
+	metaKeyResolvedProviderID = "llm.resolved_provider_id"
+	metaKeySelectedPolicyID   = "llm.selected_policy_id"
+	metaKeyPolicyDecision     = "llm_policy.decision"
+	metaKeyPolicyReason       = "llm_policy.reason"
+	metaKeyInputTokens        = "llm.input_tokens"  //nolint:gosec // metadata key name, not a credential
+	metaKeyOutputTokens       = "llm.output_tokens" //nolint:gosec // metadata key name, not a credential
+	metaKeyTotalTokens        = "llm.total_tokens"  //nolint:gosec // metadata key name, not a credential
+	metaKeyCostUSDTotal       = "cost.usd_total"
+	metaKeyStream             = "llm.stream"
+	metaKeySessionID          = "llm.session_id"
+	metaKeyAuthorisingGroups  = "llm.authorising_groups"
+	metaKeyRequestPrompt      = "llm.request_prompt"
+	metaKeyResponseCompletion = "llm.response_completion"
+)
+
 type managerImpl struct {
 	store              store.Store
 	permissionsManager permissions.Manager
@@ -31,8 +56,14 @@ func NewManager(store store.Store, permissionsManager permissions.Manager, geo g
 	}
 }

-// SaveAccessLog saves an access log entry to the database after enriching it
+// SaveAccessLog saves an access log entry to the database after enriching it.
+// Agent-network entries are flattened into their own dedicated table (queryable
+// LLM columns + group child rows) instead of the shared reverse-proxy table.
 func (m *managerImpl) SaveAccessLog(ctx context.Context, logEntry *accesslogs.AccessLogEntry) error {
+	if logEntry.AgentNetwork {
+		return m.saveAgentNetworkAccessLog(ctx, logEntry)
+	}
+
 	if m.geo != nil && logEntry.GeoLocation.ConnectionIP != nil {
 		location, err := m.geo.Lookup(logEntry.GeoLocation.ConnectionIP)
 		if err != nil {
@@ -61,6 +92,184 @@ func (m *managerImpl) SaveAccessLog(ctx context.Context, logEntry *accesslogs.Ac
 	return nil
 }

+// saveAgentNetworkAccessLog flattens the metadata-bearing access-log entry and
+// persists it in two parts:
+//
+//   - The stripped usage record is written unconditionally — usage/cost is
+//     collected on every request regardless of the account's log-collection
+//     toggle (the proxy ships a usage-only entry when logging is disabled).
+//   - The full access-log row (with request detail + prompt) is written only
+//     when the account's EnableLogCollection setting is on. This setting read
+//     is the authoritative gate; the proxy-side strip is defense in depth.
+func (m *managerImpl) saveAgentNetworkAccessLog(ctx context.Context, logEntry *accesslogs.AccessLogEntry) error {
+	entry, groups := flattenAgentNetworkLog(logEntry)
+
+	usage, usageGroups := usageFromFlattenedLog(entry, groups)
+	if err := m.store.CreateAgentNetworkUsage(ctx, usage, usageGroups); err != nil {
+		log.WithContext(ctx).WithFields(log.Fields{
+			"account_id": entry.AccountID,
+			"model":      entry.Model,
+		}).Errorf("failed to save agent-network usage: %v", err)
+		return err
+	}
+
+	settings, err := m.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, entry.AccountID)
+	if err != nil {
+		// No settings row (or a transient read error) means we can't confirm
+		// log collection is enabled — usage is already saved, so skip the full
+		// row rather than fail the whole ingest.
+		log.WithContext(ctx).Debugf("skipping full agent-network access-log row for account %s: %v", entry.AccountID, err)
+		return nil
+	}
+	if !settings.EnableLogCollection {
+		return nil
+	}
+
+	if err := m.store.CreateAgentNetworkAccessLog(ctx, entry, groups); err != nil {
+		log.WithContext(ctx).WithFields(log.Fields{
+			"account_id": entry.AccountID,
+			"service_id": entry.ServiceID,
+			"model":      entry.Model,
+			"status":     entry.StatusCode,
+		}).Errorf("failed to save agent-network access log: %v", err)
+		return err
+	}
+	return nil
+}
+
+// flattenAgentNetworkLog converts a reverse-proxy AccessLogEntry (whose LLM
+// dimensions live in the opaque Metadata map) into the flattened
+// agent-network row + authorising-group child rows.
+func flattenAgentNetworkLog(e *accesslogs.AccessLogEntry) (*agentNetworkTypes.AgentNetworkAccessLog, []agentNetworkTypes.AgentNetworkAccessLogGroup) {
+	meta := e.Metadata
+
+	var sourceIP string
+	if e.GeoLocation.ConnectionIP != nil {
+		sourceIP = e.GeoLocation.ConnectionIP.String()
+	}
+
+	entry := &agentNetworkTypes.AgentNetworkAccessLog{
+		ID:            e.ID,
+		AccountID:     e.AccountID,
+		ServiceID:     e.ServiceID,
+		Timestamp:     e.Timestamp,
+		UserID:        e.UserId,
+		SourceIP:      sourceIP,
+		Method:        e.Method,
+		Host:          e.Host,
+		Path:          e.Path,
+		Duration:      e.Duration,
+		StatusCode:    e.StatusCode,
+		AuthMethod:    e.AuthMethodUsed,
+		BytesUpload:   e.BytesUpload,
+		BytesDownload: e.BytesDownload,
+
+		Provider:           meta[metaKeyProvider],
+		Model:              meta[metaKeyModel],
+		SessionID:          meta[metaKeySessionID],
+		ResolvedProviderID: meta[metaKeyResolvedProviderID],
+		SelectedPolicyID:   meta[metaKeySelectedPolicyID],
+		Decision:           meta[metaKeyPolicyDecision],
+		DenyReason:         meta[metaKeyPolicyReason],
+		InputTokens:        parseMetaInt(meta, metaKeyInputTokens),
+		OutputTokens:       parseMetaInt(meta, metaKeyOutputTokens),
+		TotalTokens:        parseMetaInt(meta, metaKeyTotalTokens),
+		CostUSD:            parseMetaFloat(meta, metaKeyCostUSDTotal),
+		Stream:             parseMetaBool(meta, metaKeyStream),
+		RequestPrompt:      meta[metaKeyRequestPrompt],
+		ResponseCompletion: meta[metaKeyResponseCompletion],
+	}
+
+	var groups []agentNetworkTypes.AgentNetworkAccessLogGroup
+	for _, gid := range parseGroupCSV(meta[metaKeyAuthorisingGroups]) {
+		groups = append(groups, agentNetworkTypes.AgentNetworkAccessLogGroup{
+			LogID:     entry.ID,
+			GroupID:   gid,
+			AccountID: entry.AccountID,
+		})
+	}
+	return entry, groups
+}
+
+// usageFromFlattenedLog derives the stripped usage record (and its group child
+// rows) from an already-flattened access-log entry. The usage row shares the
+// log's ID so the two correlate.
+func usageFromFlattenedLog(e *agentNetworkTypes.AgentNetworkAccessLog, groups []agentNetworkTypes.AgentNetworkAccessLogGroup) (*agentNetworkTypes.AgentNetworkUsage, []agentNetworkTypes.AgentNetworkUsageGroup) {
+	usage := &agentNetworkTypes.AgentNetworkUsage{
+		ID:                 e.ID,
+		AccountID:          e.AccountID,
+		Timestamp:          e.Timestamp,
+		UserID:             e.UserID,
+		ResolvedProviderID: e.ResolvedProviderID,
+		Provider:           e.Provider,
+		Model:              e.Model,
+		SessionID:          e.SessionID,
+		InputTokens:        e.InputTokens,
+		OutputTokens:       e.OutputTokens,
+		TotalTokens:        e.TotalTokens,
+		CostUSD:            e.CostUSD,
+	}
+
+	usageGroups := make([]agentNetworkTypes.AgentNetworkUsageGroup, 0, len(groups))
+	for _, g := range groups {
+		usageGroups = append(usageGroups, agentNetworkTypes.AgentNetworkUsageGroup{
+			UsageID:   usage.ID,
+			GroupID:   g.GroupID,
+			AccountID: g.AccountID,
+		})
+	}
+	return usage, usageGroups
+}
+
+// parseMetaInt parses a non-negative token count. Negative or unparseable
+// values are clamped to 0 so a malformed metric can't persist a negative
+// counter.
+func parseMetaInt(meta map[string]string, key string) int64 {
+	if v, err := strconv.ParseInt(strings.TrimSpace(meta[key]), 10, 64); err == nil && v >= 0 {
+		return v
+	}
+	return 0
+}
+
+// parseMetaFloat parses a non-negative, finite cost. Negative, NaN, Inf, or
+// unparseable values are clamped to 0 so a malformed metric can't poison the
+// stored cost.
+func parseMetaFloat(meta map[string]string, key string) float64 {
+	if v, err := strconv.ParseFloat(strings.TrimSpace(meta[key]), 64); err == nil && v >= 0 && !math.IsInf(v, 0) {
+		return v
+	}
+	return 0
+}
+
+func parseMetaBool(meta map[string]string, key string) bool {
+	v, _ := strconv.ParseBool(strings.TrimSpace(meta[key]))
+	return v
+}
+
+// parseGroupCSV splits the comma-separated authorising-group id list the proxy
+// emits, trimming blanks and de-duplicating. Dedup matters because the group
+// rows are keyed by (log_id, group_id) / (usage_id, group_id): a repeated id
+// in the CSV would otherwise produce a duplicate primary key and fail the
+// insert transaction.
+func parseGroupCSV(raw string) []string {
+	if raw == "" {
+		return nil
+	}
+	parts := strings.Split(raw, ",")
+	out := make([]string, 0, len(parts))
+	seen := make(map[string]struct{}, len(parts))
+	for _, p := range parts {
+		if p = strings.TrimSpace(p); p != "" {
+			if _, dup := seen[p]; dup {
+				continue
+			}
+			seen[p] = struct{}{}
+			out = append(out, p)
+		}
+	}
+	return out
+}
+
 // GetAllAccessLogs retrieves access logs for an account with pagination and filtering
 func (m *managerImpl) GetAllAccessLogs(ctx context.Context, accountID, userID string, filter *accesslogs.AccessLogFilter) ([]*accesslogs.AccessLogEntry, int64, error) {
 	ok, ctx, err := m.permissionsManager.ValidateUserPermissions(ctx, accountID, userID, modules.Services, operations.Read)