mirror of
https://github.com/netbirdio/netbird.git
synced 2026-10-28 15:29:09 +01:00
[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.
This commit is contained in:
mlsmaycon
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user