[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.
This commit is contained in:
Maycon Santos
2026-06-27 13:41:00 +02:00
committed by GitHub
parent 615631567a
commit b416063bcc
187 changed files with 36835 additions and 660 deletions
@@ -0,0 +1,38 @@
// Package llm_limit_check is the SlotOnRequest middleware that asks
// management which agent-network policy "pays" for the current LLM
// request. On allow, it stamps the selected policy id, attribution
// group id, and effective window length onto the metadata bag so the
// post-flight llm_limit_record middleware can tick the right counters.
// On deny, it returns a 403 carrying the canonical llm_policy.* deny
// code surfaced by management.
package llm_limit_check
import (
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
)
// ID is the registry identifier for this middleware.
const ID = "llm_limit_check"
// Factory builds a configured llm_limit_check instance. The factory
// has no per-target config — it pulls the management gRPC client from
// the package-level FactoryContext at construction time. A nil
// MgmtClient on the context is allowed; the middleware then becomes
// a no-op pass-through (allow without attribution) so a partially
// wired environment doesn't break the chain.
type Factory struct{}
// ID returns the registry identifier matching the middleware ID.
func (Factory) ID() string { return ID }
// New ignores the rawConfig payload (no per-target config today) and
// returns a Middleware bound to the FactoryContext's MgmtClient.
func (Factory) New(_ []byte) (middleware.Middleware, error) {
ctx := builtin.Context()
return New(ctx.MgmtClient, ctx.Logger), nil
}
func init() {
builtin.Register(Factory{})
}
@@ -0,0 +1,196 @@
package llm_limit_check
import (
"context"
"strconv"
"time"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
"github.com/netbirdio/netbird/shared/management/proto"
)
// Version is reported via Middleware.Version().
const Version = "1.0.0"
// callTimeout caps the wall-clock budget for the pre-flight RPC. The
// middleware sits on the request leg, so a slow management call
// translates directly to user-visible latency. 2s is loose enough for
// a healthy management cluster but tight enough that a stalled call
// fails open via the same path nil-MgmtClient does — an enforcement
// gate that adds 30s of latency is worse than a stale gate.
const callTimeout = 2 * time.Second
// Middleware is the per-target instance that runs the pre-flight check.
type Middleware struct {
mgmt builtin.MgmtClient
logger *log.Logger
}
// New constructs a Middleware. mgmt may be nil — that's the
// no-management-wired case where the middleware is a pass-through
// (allow without attribution); useful for unit tests and for
// progressive rollout of the management RPC.
func New(mgmt builtin.MgmtClient, logger *log.Logger) *Middleware {
if logger == nil {
logger = log.StandardLogger()
}
return &Middleware{mgmt: mgmt, logger: logger}
}
// ID returns the registry identifier.
func (m *Middleware) ID() string { return ID }
// Version returns the implementation version.
func (m *Middleware) Version() string { return Version }
// Slot reports the chain slot the middleware lives in.
func (m *Middleware) Slot() middleware.Slot { return middleware.SlotOnRequest }
// AcceptedContentTypes returns nil because the gate consults metadata
// emitted upstream (KeyLLMResolvedProviderID) and never inspects bodies.
func (m *Middleware) AcceptedContentTypes() []string { return nil }
// MetadataKeys is the closed allowlist of keys this middleware emits.
func (m *Middleware) MetadataKeys() []string {
return []string{
middleware.KeyLLMSelectedPolicyID,
middleware.KeyLLMAttributionGroupID,
middleware.KeyLLMAttributionWindowS,
middleware.KeyLLMPolicyDecision,
middleware.KeyLLMPolicyReason,
}
}
// MutationsSupported reports that the middleware never mutates the
// request body or headers; the only outcome is allow + metadata or
// deny.
func (m *Middleware) MutationsSupported() bool { return false }
// Close releases resources owned by the middleware. Stateless, so
// this is a no-op.
func (m *Middleware) Close() error { return nil }
// Invoke runs the pre-flight policy check.
func (m *Middleware) Invoke(ctx context.Context, in *middleware.Input) (*middleware.Output, error) {
if m.mgmt == nil {
// No management client wired — fall through to allow with
// no attribution. RecordLLMUsage on the response leg will
// also be a no-op so counters stay at zero. This matches
// the PR1 behaviour exactly so a partial wiring is
// indistinguishable from "no enforcement".
return allowNoAttribution(), nil
}
providerID := lookupKV(in.Metadata, middleware.KeyLLMResolvedProviderID)
if providerID == "" {
// llm_router didn't emit a resolved provider id — usually
// because the request didn't carry an llm.model. The
// router itself denied; we won't reach here in production,
// but defensively pass through so we never deny on top of
// an upstream allow.
return allowNoAttribution(), nil
}
rpcCtx, cancel := context.WithTimeout(ctx, callTimeout)
defer cancel()
resp, err := m.mgmt.CheckLLMPolicyLimits(rpcCtx, &proto.CheckLLMPolicyLimitsRequest{
AccountId: in.AccountID,
UserId: in.UserID,
GroupIds: append([]string(nil), in.UserGroups...),
ProviderId: providerID,
Model: lookupKV(in.Metadata, middleware.KeyLLMModel),
})
if err != nil {
// Fail-open on transport / management errors. The
// alternative — denying every request when management is
// unreachable — is worse for v1 (operational outage =
// total LLM outage). Operators can audit via the
// access-log; PR3 can switch to fail-closed under a flag.
m.logger.WithError(err).
WithField("middleware", ID).
Debugf("management pre-flight failed; failing open")
return allowNoAttribution(), nil
}
if resp.GetDecision() == "deny" {
return denyFromManagement(resp), nil
}
return allowFromManagement(resp), nil
}
// allowNoAttribution returns the no-op allow envelope used when no
// management client is wired or no provider was resolved. Stamps
// decision=allow but no policy / attribution metadata so
// llm_limit_record skips its post-flight write.
func allowNoAttribution() *middleware.Output {
return &middleware.Output{
Decision: middleware.DecisionAllow,
Metadata: []middleware.KV{
{Key: middleware.KeyLLMPolicyDecision, Value: "allow"},
},
}
}
// allowFromManagement converts a successful CheckLLMPolicyLimits
// response into the chain's allow envelope, stamping the attribution
// metadata the response leg consumes.
func allowFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Output {
out := &middleware.Output{
Decision: middleware.DecisionAllow,
Metadata: []middleware.KV{
{Key: middleware.KeyLLMPolicyDecision, Value: "allow"},
},
}
if id := resp.GetSelectedPolicyId(); id != "" {
out.Metadata = append(out.Metadata, middleware.KV{Key: middleware.KeyLLMSelectedPolicyID, Value: id})
}
if g := resp.GetAttributionGroupId(); g != "" {
out.Metadata = append(out.Metadata, middleware.KV{Key: middleware.KeyLLMAttributionGroupID, Value: g})
}
if w := resp.GetWindowSeconds(); w > 0 {
out.Metadata = append(out.Metadata, middleware.KV{Key: middleware.KeyLLMAttributionWindowS, Value: strconv.FormatInt(w, 10)})
}
return out
}
// denyFromManagement converts a deny response into the chain's deny
// envelope. The deny code surfaces verbatim through the framework's
// fixed JSON template; arbitrary middleware bytes can't reach the
// wire.
func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Output {
code := resp.GetDenyCode()
if code == "" {
code = "llm_policy.cap_exceeded"
}
// The canonical code is safe to surface; the management-supplied
// reason can name internal quota details (used amounts, caps, rule
// ids), so keep the public message generic and leave the detail to
// server-side logs.
return &middleware.Output{
Decision: middleware.DecisionDeny,
DenyStatus: 403,
DenyReason: &middleware.DenyReason{
Code: code,
Message: "LLM policy limit exceeded",
},
Metadata: []middleware.KV{
{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},
{Key: middleware.KeyLLMPolicyReason, Value: code},
},
}
}
// lookupKV returns the value associated with key, or the empty
// string when absent.
func lookupKV(kvs []middleware.KV, key string) string {
for _, kv := range kvs {
if kv.Key == key {
return kv.Value
}
}
return ""
}
@@ -0,0 +1,186 @@
package llm_limit_check
import (
"context"
"errors"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/shared/management/proto"
)
// fakeMgmt is a minimal builtin.MgmtClient stub that lets the test
// drive CheckLLMPolicyLimits responses without a real gRPC dial.
type fakeMgmt struct {
checkResp *proto.CheckLLMPolicyLimitsResponse
checkErr error
checkReq *proto.CheckLLMPolicyLimitsRequest
}
func (f *fakeMgmt) CheckLLMPolicyLimits(_ context.Context, in *proto.CheckLLMPolicyLimitsRequest, _ ...grpc.CallOption) (*proto.CheckLLMPolicyLimitsResponse, error) {
f.checkReq = in
return f.checkResp, f.checkErr
}
func (f *fakeMgmt) RecordLLMUsage(_ context.Context, _ *proto.RecordLLMUsageRequest, _ ...grpc.CallOption) (*proto.RecordLLMUsageResponse, error) {
return &proto.RecordLLMUsageResponse{}, nil
}
func runInvoke(t *testing.T, m *Middleware, in *middleware.Input) *middleware.Output {
t.Helper()
out, err := m.Invoke(context.Background(), in)
require.NoError(t, err, "Invoke must not propagate transport errors")
require.NotNil(t, out, "Invoke must always return an Output")
return out
}
// TestInvoke_AllowStampsAttributionMetadata covers the happy path:
// management returns an allow decision with selected_policy_id +
// attribution_group_id + window_seconds, the middleware emits all three
// onto the metadata bag so the post-flight llm_limit_record
// middleware has everything it needs to tick the right counter.
func TestInvoke_AllowStampsAttributionMetadata(t *testing.T) {
mgmt := &fakeMgmt{
checkResp: &proto.CheckLLMPolicyLimitsResponse{
Decision: "allow",
SelectedPolicyId: "pol-X",
AttributionGroupId: "grp-engineers",
WindowSeconds: 86_400,
},
}
m := New(mgmt, nil)
out := runInvoke(t, m, &middleware.Input{
AccountID: "acc-1",
UserID: "user-bob",
UserGroups: []string{"grp-engineers"},
Metadata: []middleware.KV{
{Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"},
{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
},
})
assert.Equal(t, middleware.DecisionAllow, out.Decision)
assert.Equal(t, "acc-1", mgmt.checkReq.GetAccountId(), "account_id must round-trip onto the RPC")
assert.Equal(t, "user-bob", mgmt.checkReq.GetUserId())
assert.Equal(t, []string{"grp-engineers"}, mgmt.checkReq.GetGroupIds())
assert.Equal(t, "prov-1", mgmt.checkReq.GetProviderId(), "resolved provider id must come from metadata")
assert.Equal(t, "gpt-4o", mgmt.checkReq.GetModel(), "model must come from metadata")
want := map[string]string{
middleware.KeyLLMPolicyDecision: "allow",
middleware.KeyLLMSelectedPolicyID: "pol-X",
middleware.KeyLLMAttributionGroupID: "grp-engineers",
middleware.KeyLLMAttributionWindowS: "86400",
}
got := map[string]string{}
for _, kv := range out.Metadata {
got[kv.Key] = kv.Value
}
assert.Equal(t, want, got, "attribution metadata must land on the bag for the response leg to consume")
}
// TestInvoke_DenyConvertsToProxyDeny proves the deny envelope round-
// trips: management's deny code becomes the proxy framework's deny
// payload at status 403, and the deny reason text is preserved so
// operators can debug from the access log.
func TestInvoke_DenyConvertsToProxyDeny(t *testing.T) {
mgmt := &fakeMgmt{
checkResp: &proto.CheckLLMPolicyLimitsResponse{
Decision: "deny",
DenyCode: "llm_policy.token_cap_exceeded",
DenyReason: "group token cap exhausted on policy pol-X (used 1000 of 1000)",
},
}
m := New(mgmt, nil)
out := runInvoke(t, m, &middleware.Input{
AccountID: "acc-1",
UserGroups: []string{"grp-engineers"},
Metadata: []middleware.KV{{Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"}},
})
assert.Equal(t, middleware.DecisionDeny, out.Decision)
assert.Equal(t, 403, out.DenyStatus, "policy denials are 403 — same as llm_router's")
require.NotNil(t, out.DenyReason, "deny envelope must carry a reason payload")
assert.Equal(t, "llm_policy.token_cap_exceeded", out.DenyReason.Code, "canonical deny code surfaces to the caller")
// The public message must stay generic: the management reason names
// internal quota detail (used/cap, rule id) that must not leak.
assert.Equal(t, "LLM policy limit exceeded", out.DenyReason.Message, "public deny message must be generic")
assert.NotContains(t, out.DenyReason.Message, "exhausted", "internal quota detail must not reach the caller")
assert.NotContains(t, out.DenyReason.Message, "1000", "internal cap numbers must not reach the caller")
}
// TestInvoke_NoMgmtClientPassesThrough proves the partial-wiring
// safety: a middleware constructed without a management client
// allows every request without attribution. This makes a half-set-up
// environment indistinguishable from "no enforcement" rather than
// breaking the chain.
func TestInvoke_NoMgmtClientPassesThrough(t *testing.T) {
m := New(nil, nil)
out := runInvoke(t, m, &middleware.Input{
AccountID: "acc-1",
UserGroups: []string{"grp-engineers"},
Metadata: []middleware.KV{{Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"}},
})
assert.Equal(t, middleware.DecisionAllow, out.Decision)
for _, kv := range out.Metadata {
assert.NotEqual(t, middleware.KeyLLMSelectedPolicyID, kv.Key,
"no mgmt client = no attribution metadata; record middleware then skips its write")
}
}
// TestInvoke_NoResolvedProviderPassesThrough covers the defensive
// path: when llm_router didn't set llm.resolved_provider_id (which
// only happens on the deny side of llm_router), the gate must NOT
// stack a second deny on top — pass through and let the upstream
// deny stand.
func TestInvoke_NoResolvedProviderPassesThrough(t *testing.T) {
m := New(&fakeMgmt{}, nil)
out := runInvoke(t, m, &middleware.Input{
AccountID: "acc-1",
Metadata: []middleware.KV{},
})
assert.Equal(t, middleware.DecisionAllow, out.Decision,
"no resolved provider = the gate has nothing to check; never deny on top of an upstream allow")
}
// TestInvoke_RPCErrorFailsOpen proves the fail-open contract: a
// transport error from management does NOT deny the request. v1
// trades enforcement strictness for availability — an unreachable
// management server otherwise turns into a total LLM outage.
func TestInvoke_RPCErrorFailsOpen(t *testing.T) {
m := New(&fakeMgmt{checkErr: errors.New("connection refused")}, nil)
out := runInvoke(t, m, &middleware.Input{
AccountID: "acc-1",
UserGroups: []string{"grp-engineers"},
Metadata: []middleware.KV{{Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"}},
})
assert.Equal(t, middleware.DecisionAllow, out.Decision,
"transport errors must not cascade into total LLM outages — operators audit via access log")
}
// TestMetadataKeys_Allowlist locks the closed set this middleware can
// emit. The accumulator drops anything outside this list; adding a
// new emission means updating both the slice and this test.
func TestMetadataKeys_Allowlist(t *testing.T) {
keys := New(nil, nil).MetadataKeys()
want := []string{
middleware.KeyLLMSelectedPolicyID,
middleware.KeyLLMAttributionGroupID,
middleware.KeyLLMAttributionWindowS,
middleware.KeyLLMPolicyDecision,
middleware.KeyLLMPolicyReason,
}
assert.ElementsMatch(t, want, keys)
}