mirror of
https://github.com/netbirdio/netbird.git
synced 2026-10-04 04:29:06 +02:00
[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
@@ -0,0 +1,320 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// boundMiddleware pairs a validated spec with the resolved middleware
|
||||
// instance the chain will invoke.
|
||||
type boundMiddleware struct {
|
||||
spec Spec
|
||||
mw Middleware
|
||||
}
|
||||
|
||||
// Chain is the ordered set of middlewares that run for a specific
|
||||
// target. Chains are immutable once built; Manager produces a new
|
||||
// Chain on every Rebuild.
|
||||
//
|
||||
// Ordering: middlewares are kept in registration order. RunRequest
|
||||
// iterates the SlotOnRequest middlewares in order; RunResponse
|
||||
// iterates the SlotOnResponse middlewares in reverse order
|
||||
// (middleware-style LIFO so the last to see the request is the first
|
||||
// to see the response); RunTerminal iterates the SlotTerminal
|
||||
// middlewares in registration order, after every on_response slot has
|
||||
// emitted, so the metadata bag they observe is complete.
|
||||
//
|
||||
// Close drains in-flight invocations and tears down each middleware.
|
||||
// Callers swapping a chain via Manager invoke Close on the old chain
|
||||
// after the swap so live requests finish on the previous instance.
|
||||
type Chain struct {
|
||||
targetID string
|
||||
all []boundMiddleware
|
||||
onRequest []int
|
||||
onResponse []int
|
||||
terminal []int
|
||||
dispatcher *Dispatcher
|
||||
inflight sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewChain assembles a Chain from the bound middlewares. The slice
|
||||
// order is the registration order; the chain captures index slices
|
||||
// per slot so iteration does not re-scan the slot field per call.
|
||||
func NewChain(targetID string, bound []boundMiddleware, d *Dispatcher) *Chain {
|
||||
c := &Chain{
|
||||
targetID: targetID,
|
||||
all: bound,
|
||||
dispatcher: d,
|
||||
}
|
||||
for i, bm := range bound {
|
||||
switch bm.spec.Slot {
|
||||
case SlotOnRequest:
|
||||
c.onRequest = append(c.onRequest, i)
|
||||
case SlotOnResponse:
|
||||
c.onResponse = append(c.onResponse, i)
|
||||
case SlotTerminal:
|
||||
c.terminal = append(c.terminal, i)
|
||||
}
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// Close waits for outstanding invocations against this chain to
|
||||
// finish (bounded by ctx) and releases the middleware instances bound
|
||||
// to it. Safe to call once the chain has been removed from the
|
||||
// routing snapshot. Subsequent Run* calls are still safe (return
|
||||
// without invoking) but Close itself is one-shot.
|
||||
func (c *Chain) Close(ctx context.Context) error {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
c.inflight.Wait()
|
||||
close(done)
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
case <-ctx.Done():
|
||||
// Drain timed out: requests may still be running against these
|
||||
// middleware instances, so tearing them down now risks a
|
||||
// use-after-close. Leave them (a bounded leak) and surface the
|
||||
// timeout; the runaway backstop in the Manager already alerts.
|
||||
return ctx.Err()
|
||||
}
|
||||
for _, bm := range c.all {
|
||||
if bm.mw == nil {
|
||||
continue
|
||||
}
|
||||
if err := bm.mw.Close(); err != nil {
|
||||
c.dispatcher.logger.Debugf("middleware %s close: %v", bm.spec.ID, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Empty reports whether the chain has no middlewares.
|
||||
func (c *Chain) Empty() bool {
|
||||
return c == nil || len(c.all) == 0
|
||||
}
|
||||
|
||||
// TargetID returns the key used to find this chain.
|
||||
func (c *Chain) TargetID() string {
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
return c.targetID
|
||||
}
|
||||
|
||||
// IDs returns the ordered list of middleware IDs bound to this chain.
|
||||
func (c *Chain) IDs() []string {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
out := make([]string, len(c.all))
|
||||
for i, bm := range c.all {
|
||||
out[i] = bm.spec.ID
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// RunRequest iterates the on_request slot in registration order. Deny
|
||||
// short-circuits the remaining middlewares and returns the deny
|
||||
// output. The caller owns applying mutations to the real request and
|
||||
// merging the metadata returned in `merged` into the captured-data
|
||||
// bag passed to subsequent slots.
|
||||
//
|
||||
// Each middleware sees the metadata emitted by earlier middlewares in
|
||||
// the same slot — this is how llm_guardrail reads
|
||||
// llm.request_prompt_raw from llm_request_parser without a side
|
||||
// channel, and how cost_meter reads tokens emitted by
|
||||
// llm_response_parser on the response leg.
|
||||
//
|
||||
// If any middleware emits a non-nil Mutations.RewriteUpstream while
|
||||
// satisfying the mutation gates (CanMutate && MutationsSupported), the
|
||||
// latest such value is returned to the caller. Last-write-wins so the
|
||||
// last middleware in the slot can override an earlier rewrite.
|
||||
func (c *Chain) RunRequest(ctx context.Context, r *http.Request, in *Input, acc *Accumulator) (denied *Output, merged []KV, rewrite *UpstreamRewrite, err error) {
|
||||
if c.Empty() || len(c.onRequest) == 0 {
|
||||
return nil, nil, nil, nil
|
||||
}
|
||||
c.inflight.Add(1)
|
||||
defer c.inflight.Done()
|
||||
running := append([]KV(nil), in.Metadata...)
|
||||
for _, idx := range c.onRequest {
|
||||
bm := c.all[idx]
|
||||
call := cloneInputFor(in, SlotOnRequest)
|
||||
call.Metadata = append([]KV(nil), running...)
|
||||
out, invErr := c.dispatcher.Invoke(ctx, bm.spec, bm.mw, call)
|
||||
if invErr != nil && out == nil {
|
||||
continue
|
||||
}
|
||||
if out == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
accepted, rejected := acc.Emit(bm.spec.ID, bm.spec.MetadataKeys, out.Metadata)
|
||||
for _, rej := range rejected {
|
||||
c.dispatcher.metrics.IncMetadataRejected(ctx, bm.spec.ID, rej.Reason)
|
||||
}
|
||||
merged = append(merged, accepted...)
|
||||
running = append(running, accepted...)
|
||||
|
||||
if out.Decision == DecisionDeny {
|
||||
c.dispatcher.metrics.IncRequest(ctx, bm.spec.ID, c.targetID, "deny")
|
||||
return out, merged, rewrite, nil
|
||||
}
|
||||
c.dispatcher.metrics.IncRequest(ctx, bm.spec.ID, c.targetID, "allow")
|
||||
|
||||
if rw := mutationRewrite(bm.spec, out.Mutations); rw != nil {
|
||||
rewrite = rw
|
||||
}
|
||||
if r != nil && bm.spec.CanMutate && out.Mutations != nil {
|
||||
applyMutations(ctx, c.dispatcher, bm.spec, r, out.Mutations)
|
||||
}
|
||||
}
|
||||
return nil, merged, rewrite, nil
|
||||
}
|
||||
|
||||
// RunResponse iterates the on_response slot in reverse registration
|
||||
// order, matching the middleware "last in, first out" convention so
|
||||
// the last middleware to see the request is the first to see the
|
||||
// response. Middlewares cannot deny; they emit metadata.
|
||||
//
|
||||
// As with RunRequest, each middleware sees the metadata emitted by
|
||||
// earlier middlewares in this slot — accumulated in the order the
|
||||
// middlewares run (LIFO of registration). cost_meter relies on this
|
||||
// to read llm.input_tokens / llm.output_tokens that
|
||||
// llm_response_parser emitted just before it.
|
||||
func (c *Chain) RunResponse(ctx context.Context, in *Input, acc *Accumulator) (merged []KV) {
|
||||
if c.Empty() || len(c.onResponse) == 0 {
|
||||
return nil
|
||||
}
|
||||
c.inflight.Add(1)
|
||||
defer c.inflight.Done()
|
||||
running := append([]KV(nil), in.Metadata...)
|
||||
for i := len(c.onResponse) - 1; i >= 0; i-- {
|
||||
bm := c.all[c.onResponse[i]]
|
||||
call := cloneInputFor(in, SlotOnResponse)
|
||||
call.Metadata = append([]KV(nil), running...)
|
||||
out, _ := c.dispatcher.Invoke(ctx, bm.spec, bm.mw, call)
|
||||
if out == nil {
|
||||
continue
|
||||
}
|
||||
accepted, rejected := acc.Emit(bm.spec.ID, bm.spec.MetadataKeys, out.Metadata)
|
||||
for _, rej := range rejected {
|
||||
c.dispatcher.metrics.IncMetadataRejected(ctx, bm.spec.ID, rej.Reason)
|
||||
}
|
||||
merged = append(merged, accepted...)
|
||||
running = append(running, accepted...)
|
||||
c.dispatcher.metrics.IncRequest(ctx, bm.spec.ID, c.targetID, "passthrough")
|
||||
}
|
||||
return merged
|
||||
}
|
||||
|
||||
// RunTerminal iterates the terminal slot in registration order, after
|
||||
// every on_response middleware has emitted. Terminal middlewares
|
||||
// observe the full metadata bag carried in `in.Metadata` plus any
|
||||
// emissions from terminal middlewares that ran before them; they
|
||||
// cannot deny and cannot mutate.
|
||||
func (c *Chain) RunTerminal(ctx context.Context, in *Input, acc *Accumulator) (merged []KV) {
|
||||
if c.Empty() || len(c.terminal) == 0 {
|
||||
return nil
|
||||
}
|
||||
c.inflight.Add(1)
|
||||
defer c.inflight.Done()
|
||||
running := append([]KV(nil), in.Metadata...)
|
||||
for _, idx := range c.terminal {
|
||||
bm := c.all[idx]
|
||||
call := cloneInputFor(in, SlotTerminal)
|
||||
call.Metadata = append([]KV(nil), running...)
|
||||
out, _ := c.dispatcher.Invoke(ctx, bm.spec, bm.mw, call)
|
||||
if out == nil {
|
||||
continue
|
||||
}
|
||||
accepted, rejected := acc.Emit(bm.spec.ID, bm.spec.MetadataKeys, out.Metadata)
|
||||
for _, rej := range rejected {
|
||||
c.dispatcher.metrics.IncMetadataRejected(ctx, bm.spec.ID, rej.Reason)
|
||||
}
|
||||
merged = append(merged, accepted...)
|
||||
running = append(running, accepted...)
|
||||
c.dispatcher.metrics.IncRequest(ctx, bm.spec.ID, c.targetID, "terminal")
|
||||
}
|
||||
return merged
|
||||
}
|
||||
|
||||
// mutationRewrite returns the upstream rewrite carried in m when the
|
||||
// spec's mutation gates allow it. The rewrite itself is not applied
|
||||
// here; the caller (reverse proxy) decides whether to honour it.
|
||||
func mutationRewrite(spec Spec, m *Mutations) *UpstreamRewrite {
|
||||
if m == nil || m.RewriteUpstream == nil {
|
||||
return nil
|
||||
}
|
||||
if !spec.CanMutate || !spec.MutationsSupported {
|
||||
return nil
|
||||
}
|
||||
return m.RewriteUpstream
|
||||
}
|
||||
|
||||
func applyMutations(ctx context.Context, d *Dispatcher, spec Spec, r *http.Request, m *Mutations) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
add, remove, blocked := FilterHeaderMutations(m)
|
||||
for _, h := range blocked {
|
||||
d.metrics.IncHeaderMutationBlocked(ctx, spec.ID, h)
|
||||
}
|
||||
for _, name := range remove {
|
||||
r.Header.Del(name)
|
||||
}
|
||||
for _, kv := range add {
|
||||
r.Header.Add(kv.Key, kv.Value)
|
||||
}
|
||||
if len(m.BodyReplace) == 0 {
|
||||
return
|
||||
}
|
||||
if err := ValidateBodyReplace(r, m.BodyReplace, true); err != nil {
|
||||
d.logger.Warnf("middleware %s body replace rejected: %v", spec.ID, err)
|
||||
return
|
||||
}
|
||||
ApplyBodyReplace(r, m.BodyReplace)
|
||||
}
|
||||
|
||||
// cloneInputFor deep-copies the mutation-prone fields of Input so
|
||||
// each middleware receives an isolated view.
|
||||
func cloneInputFor(in *Input, slot Slot) *Input {
|
||||
if in == nil {
|
||||
return nil
|
||||
}
|
||||
out := *in
|
||||
out.Slot = slot
|
||||
out.Headers = cloneKVs(in.Headers)
|
||||
out.RespHeaders = cloneKVs(in.RespHeaders)
|
||||
out.Metadata = cloneKVs(in.Metadata)
|
||||
if len(in.UserGroups) > 0 {
|
||||
out.UserGroups = append([]string(nil), in.UserGroups...)
|
||||
}
|
||||
if len(in.UserGroupNames) > 0 {
|
||||
out.UserGroupNames = append([]string(nil), in.UserGroupNames...)
|
||||
}
|
||||
if len(in.Body) > 0 {
|
||||
out.Body = append([]byte(nil), in.Body...)
|
||||
}
|
||||
if len(in.RespBody) > 0 {
|
||||
out.RespBody = append([]byte(nil), in.RespBody...)
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
func cloneKVs(in []KV) []KV {
|
||||
if len(in) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]KV, len(in))
|
||||
copy(out, in)
|
||||
return out
|
||||
}
|
||||
Reference in New Issue
Block a user