Back to Netbird

proxy/middleware-builtin — the LLM chain

docs/agent-networks/modules/31-proxy-middleware-builtin.md

0.77.124.3 KB
Original Source

proxy/middleware-builtin — the LLM chain

The registry-mounted middleware set the proxy executes on every agent-network LLM request. The two highest-blast-radius areas are the capture-pointer semantics and the limit_check ⇒ limit_record record-once invariant.

Sibling module: 32-proxy-llm-parsers.md — the SDK adapters + pricing table and cost formula this chain delegates to.


Module boundary

This module is the registry-mounted middleware set the proxy executes on every agent-network LLM request. Each sub-package registers itself via init() (builtin.go:32–34); the proxy server anonymous-imports the set (all_test.go:11–19) so the registry is populated at boot. The chain is wired by the management synthesiser and executed by the framework (proxy/internal/middleware/{chain,dispatcher,accumulator}.go — both out of scope). Everything here reads from / writes to one envelope: the middleware.KV metadata bag plus middleware.Mutations for header/body rewrites.

The 8 middlewares

NameSlotInputs (metadata read)Outputs (metadata written)Side effects
llm_request_parserOnRequestInput.{URL,Body,BodyTruncated}llm.{provider,model,stream,request_prompt_raw,capture_truncated}none
llm_routerOnRequestllm.model, Input.{URL,UserGroups}llm.{resolved_provider_id,authorising_groups}, llm_policy.{decision,reason}upstream rewrite + auth strip/inject
llm_limit_checkOnRequestllm.{resolved_provider_id,model}, Input.{AccountID,UserID,UserGroups}llm.{selected_policy_id,attribution_group_id,attribution_window_seconds}, llm_policy.{decision,reason}gRPC CheckLLMPolicyLimits
llm_identity_injectOnRequestllm.{resolved_provider_id,authorising_groups}, Input.{UserEmail,UserID,UserGroups,UserGroupNames}noneheader strip/inject + optional body rewrite
llm_guardrailOnRequestllm.{model,request_prompt_raw}llm_policy.{decision,reason}, llm.request_promptnone (model allowlist deny)
llm_response_parserOnResponsellm.provider, Input.{RespHeaders,RespBody,Status}llm.{input,output,total,cached_input,cache_creation}_tokens, llm.response_completionnone
cost_meterOnResponsellm.{provider,model,resolved_provider_id}, token bucketscost.usd_{input,cached_input,cache_creation,output,total,cache} or cost.skippednone (in-memory pricing lookup)
llm_limit_recordOnResponsellm.{attribution_group_id,attribution_window_seconds,input_tokens,output_tokens}, cost.usd_totalnonegRPC RecordLLMUsage

all_test.go:26–40 locks the ID set; adding or removing one is a conscious extension.

Files

FileLOCNotes
builtin.go90Registry + FactoryContext (ctx, meter, logger, mgmt client)
all_test.go41Locks the 8-ID registry surface
agentnetwork_chain_integration_test.go319Live sqlite + real gRPC bufconn; gate→recorder wire path
llm_request_parser/*162 / 66 / 356Provider detection, body parse, prompt extraction with capture-pointer gating
llm_router/*385 / 84 / 586Three-pass route selection (model → groups → path-prefix)
llm_limit_check/*196 / 38 / 182Pre-flight CheckLLMPolicyLimits (2s, fail-open)
llm_identity_inject/*440 / 108 / 666HeaderPair (LiteLLM) + JSONMetadata (Portkey) + ExtraHeaders
llm_guardrail/*176 / 82 / 75 / 219 / 217Model allowlist + optional prompt capture with PII redaction
llm_response_parser/*258 / 222 / 43 / 433 / 169 / 111Buffered + SSE accumulation; AWS event-stream accumulator (streaming_bedrock.go) for Bedrock; capture-pointer gates completion emit
cost_meter/*236 / 98 / 586Token → USD via proxy/internal/llm/pricing; both pricing tiers arrive in the middleware config
llm_limit_record/*144 / 35 / 191Post-flight RecordLLMUsage (5s, debug-on-error)

Per-middleware

llm_request_parser

Detects the LLM provider via llm.DetectParser (URL sniff) or by name via llm.ParserByName when synthesiser stamps provider_id (middleware.go:96–99). Path-routed providers short-circuit first: parseVertexPath and parseBedrockPath (middleware.go:85–94) pull the model + vendor out of the URL before parser selection runs — Vertex from /v1/projects/.../publishers/{pub}/models/{model}:{action} (publisher → vendor via vertexPublisherVendor), Bedrock from /model/{id}/{action} with normalizeBedrockModel stripping the region prefix + version suffix. See 50-path-routed-providers.md for the full path grammar. For body-routed providers it decodes the body into RequestFacts (model + stream) and extracts the prompt. On capture_prompt=true (or absent — see capture-pointer semantics below) the prompt is run through llm_guardrail.RedactPII when redact_pii=true and truncated rune-safely to 3500 bytes (middleware.go:109–122). Key invariant: redaction is parser-side, not guardrail-side — access-log reads llm.request_prompt_raw directly.

llm_router

Three-pass route selection in matchRoute (middleware.go:241–300): filter by Models claim → vendor-pin (a vendor-tagged request never crosses to another vendor's route) → filter by AllowedGroupIDs intersection → model precedence over path → tie-break by longest UpstreamPath prefix match. Model-miss returns llm_policy.model_not_routable; known-but-unauthorised returns llm_policy.no_authorised_provider. Key invariant: auth-header strip+inject rides on UpstreamRewrite.{StripHeaders,AuthHeader} (middleware.go:606–646) — NOT HeadersAdd/HeadersRemove — because the framework's mutation gate blocks Authorization on the generic header path.

Path-routed providers route before the model table. Invoke checks isVertexPath / isBedrockPath (middleware.go:138–216) ahead of the model lookup, so a path-carried model can't be claimed by a same-vendor body-routed provider. matchPathRoute enforces the route's Models allowlist (empty = catch-all) even though the model came from the URL. Two path-only behaviours:

  • Vertex unmeterable publisher — when llm_request_parser emits no llm.provider (e.g. Gemini/google), the router denies with llm_policy.unmeterable_publisher (403) rather than forward it uncounted.
  • GCP token minting — when the route carries GCPServiceAccountKeyB64 (set from a keyfile:: api_key), gcpBearer mints + caches a short-lived OAuth2 token per request instead of injecting a static value; a bad key or unreachable token endpoint denies with llm_policy.upstream_auth_failed (502). Bedrock uses its static bearer token directly (no minting).
  • /bedrock prefix — an optional /bedrock gateway-namespace prefix is accepted and stripped via RewriteUpstream.StripPathPrefix so the native /model/... path reaches the upstream.

Full treatment in 50-path-routed-providers.md.

llm_limit_check

Pre-flight gate. Reads llm.resolved_provider_id, calls CheckLLMPolicyLimits with a 2s context timeout (middleware.go:24, 97–106), on allow stamps llm.selected_policy_id, llm.attribution_group_id, llm.attribution_window_seconds. Key invariant: fail-open. Nil MgmtClient, empty provider id, or RPC error returns allowNoAttribution() — management outage doesn't take down every LLM request. Operators audit via the access-log; a future flag may switch this to fail-closed.

llm_identity_inject

Dispatches per-rule between LiteLLM-shaped HeaderPair (middleware.go:169) and Portkey-shaped JSONMetadata (middleware.go:292). Identity is the peer's email (or UserID fallback); tags are the authorising-groups intersection emitted by llm_router, not the full UserGroups — a peer in 5 groups authorised under 1 only tags as that 1. Anti-spoof: every HeadersAdd is preceded by a HeadersRemove of the same name; the framework runs Remove before Add so client-supplied identity never reaches the upstream. Body-level inject (tags_in_body, end_user_id_in_body) is skipped on empty / truncated / non-JSON bodies so header attribution stays intact.

llm_guardrail

Model allowlist deny + optional prompt-capture-with-redaction. Allowlist match is case-insensitive via normaliseModel; empty allowlist disables the check. Prompt capture reads llm.request_prompt_raw and emits llm.request_prompt only when prompt_capture.enabled (middleware.go:149–165). Key invariant: RedactPII is the exported function the parsers call — single PII contract across all three keys.

llm_response_parser

Buffered and SSE paths share one Invoke (middleware.go:102–127): content-type sniffing dispatches to invokeBuffered (JSON, status<400) or invokeStreaming (text/event-stream, partial bodies tolerated). Streaming delegates to accumulateStream (streaming.go:21–30) using llm.NewScanner. A third path, accumulateBedrockStream (streaming_bedrock.go), decodes the AWS binary event-stream (application/vnd.amazon.eventstream) returned by Bedrock's -stream actions — InvokeModel chunk frames wrap a base64 Anthropic event, Converse frames carry text + a trailing usage block. Cached / cache-creation buckets emit only when non-zero, preserving the existing token schema.

cost_meter

Reads llm.provider + llm.model + token buckets, looks up the per-1k rates, and emits the full cost.usd_* breakdown (four per-bucket values plus the _total and _cache aggregates) or a closed-set cost.skipped reason (missing_provider/model/tokens, unparseable_tokens, zero_tokens, unknown_model).

Management owns pricing. The proxy carries no embedded price list: the whole table arrives in this middleware's ConfigJSON as {pricing: {defaults, providers}}, synthesized by management from the catalog plus the operator's stored per-provider prices (factory.go:13–34). Both tiers are validated by pricing.NewTable / pricing.NewEntries at construction, so a non-finite or negative rate fails the chain build. A price change is an ordinary mapping push — the chain rebuild yields a fresh instance over a fresh immutable table, so there is no data dir, no pricing file, no reload goroutine, and nothing to invalidate.

Two-tier lookup (middleware.go:165–183):

  1. Per-provider-record — the operator's stored price for the route that actually served the request, keyed by the llm.resolved_provider_id that llm_router stamped on the allow path, then by normalized model id. Entries arrive fully materialized (management folds default cache rates in at synth time), so there is no merging here. Absent metadata — no router in the chain — skips this tier.
  2. Surface defaults — the catalog-derived table keyed by llm.provider (openai/anthropic/bedrock). This is also what prices gateway-style providers, which enumerate no models and therefore get no per-record entry.

Backward compatibility: a config with no pricing block means management predates config-delivered pricing. The factory logs one warning at build time and the instance records cost.skipped=unknown_model ($0) for every request rather than falling back to a stale built-in price list (factory.go:55–60).

Key invariant: the provider-shape switch lives in pricing.EntryCosts (sibling doc) and is selected by the surface, not by which tier the entry came from — cost_meter stays provider-agnostic, and a per-record override on an Anthropic route still bills its cache buckets additively.

llm_limit_record

Post-flight write. Always returns DecisionAllow; response has already been served so RPC errors mustn't surface (logged at Debugf). Skip-on-no-signal at line 81 (zero tokens + zero cost). Key invariant: the skip-on-missing-attribution guard at line 98 is a safety net independent of the framework's deny short-circuit — if the gate denied and the framework still runs the recorder, the recorder skips on absent UserID+groupID+UserGroups and no phantom counter materialises.

Full-chain diagram (canonical order)

mermaid
flowchart TD
    A[HTTP request] --> B[llm_request_parser
OnRequest]
    B -->|llm.provider, llm.model,
llm.stream, llm.request_prompt_raw| C[llm_router
OnRequest]
    C -->|llm.resolved_provider_id,
llm.authorising_groups,
upstream rewrite + auth| D[llm_limit_check
OnRequest]
    D -->|deny path| Z1[403 llm_policy.*]
    D -->|allow + llm.selected_policy_id,
llm.attribution_group_id,
llm.attribution_window_seconds| E[llm_identity_inject
OnRequest]
    E -->|header strip+inject
+ optional body rewrite| F[llm_guardrail
OnRequest]
    F -->|deny: model_blocked| Z2[403 llm_policy.model_blocked]
    F -->|allow + llm.request_prompt| G[upstream LLM call]
    G --> H[llm_response_parser
OnResponse]
    H -->|llm.{input,output,total,cached_input,cache_creation}_tokens,
llm.response_completion| I[cost_meter
OnResponse]
    I -->|cost.usd_total or cost.skipped| J[llm_limit_record
OnResponse]
    J --> K[response to client]

limit_check ⇒ limit_record record-once invariant

mermaid
sequenceDiagram
    participant LC as llm_limit_check
    participant M as management gRPC
    participant U as upstream LLM
    participant LR as llm_limit_record
    participant DB as sqlite consumption table

    LC->>M: CheckLLMPolicyLimits (2s)
    alt allow
        M-->>LC: selected_policy_id, attribution_group_id, window_s
        LC->>U: stamps attribution metadata
        U-->>LR: response + tokens (via llm_response_parser + cost_meter)
        LR->>M: RecordLLMUsage (5s, debug-on-error)
        M->>DB: increment (user, group, window) row
    else deny
        M-->>LC: llm_policy.token_cap_exceeded
        Note over LR: framework short-circuits; even if invoked,
recorder skips on absent UserID+groupID+UserGroups
    else mgmt nil / rpc error
        LC-->>LC: allowNoAttribution() — fail open
        Note over LR: no window_s ⇒ recorder books only account-level
budget rules (which run independently)
    end

The integration test agentnetwork_chain_integration_test.go exercises all three branches against a real sqlite store + bufconn gRPC — no mocks. Tests: TestChain_AllowPath_StampsAttributionAndRecordsCounter (line 130), TestChain_DenyPath_GateRejectsAndNoConsumptionWritten (line 207), TestChain_CapExhaustTransition (line 265).

Public contracts (per-middleware JSON config)

MiddlewareConfig shape
llm_request_parser{provider_id?, redact_pii?, capture_prompt?: *bool} (factory.go:19–37)
llm_router{providers: [{id, models, upstream_scheme, upstream_host, upstream_path?, auth_header_name, auth_header_value, allowed_group_ids}]}
llm_limit_check{} — pulls MgmtClient from FactoryContext
llm_identity_inject`{providers: [{provider_id, header_pair?
llm_guardrail{provider_allowlists: {providerID: []string}, prompt_capture: {enabled, redact_pii}} — allowlist keyed by resolved provider id; a provider absent from the map is unrestricted (fail-closed backstop; authoritative per-policy/group check is management's CheckLLMPolicyLimits)
llm_response_parser{redact_pii?, capture_completion?: *bool}
cost_meter{pricing: {defaults: {surface: {model: rates}}, providers: {providerRecordID: {model: rates}}}} — rates are {input_per_1k, output_per_1k, cached_input_per_1k?, cache_read_per_1k?, cache_creation_per_1k?}. A missing pricing key means "management predates config-delivered pricing": every request records cost.skipped=unknown_model
llm_limit_record{} — same pattern as llm_limit_check

All factories accept empty / null / {} / whitespace as zero-value config; only structurally invalid JSON is rejected so misconfig surfaces at chain build time. cost_meter adds a semantic check on top of that: a pricing block carrying a negative or non-finite rate fails the build too, rather than mispricing live traffic.

Invariants

  1. limit_check ↔ limit_record paired. They MUST appear together. Gate stamps attribution metadata on the request leg; recorder reads it on the response leg. If a chain contains only the recorder, the skip-on-missing-attribution guard at llm_limit_record/middleware.go:81–87, 98–103 keeps counters consistent but no enforcement runs. Only-gate means counters never tick and headroom appears infinite.

  2. capture_prompt / capture_completion pointer semantics. Both are *bool. nil = "preserve legacy emit" (back-compat default for non-agent-network callers and pre-toggle tests). false = suppress the key entirely (access-log row carries zero prompt / completion content). true = emit. The synthesiser sets the pointer explicitly to the account's EnablePromptCollection toggle. The handling lives in llm_request_parser/factory.go:55–61 and the symmetric llm_response_parser/middleware.go:62–68; a missing pointer must not be treated as false (that would suppress capture for legacy non-agent-network callers). redact_pii is an orthogonal bool controlling form of emitted content, not whether it's emitted.

  3. redact_pii is parser-side. Both parsers import llm_guardrail.RedactPII and run it BEFORE stamping the metadata bag. Load-bearing because the access-log sink reads llm.request_prompt_raw and llm.response_completion directly — by the time llm_guardrail runs its own pass on llm.request_prompt, the raw key has already been stamped. Tests: TestInvoke_RedactPii_RedactsBeforeEmittingRawPrompt, TestInvoke_RedactPii_RedactsCompletionBeforeEmit.

  4. Metadata allowlist enforcement. Every middleware declares MetadataKeys(). The framework accumulator drops any KV outside that allowlist. When adding a new key, also extend the docstring in middleware/keys.go.

  5. Closed deny-code set. All deny paths emit one of: llm_policy.model_not_routable, llm_policy.no_authorised_provider, llm_policy.model_blocked, llm_policy.token_cap_exceeded, llm_policy.unmeterable_publisher (path-routed Vertex publisher with no parser → 403), llm_policy.upstream_auth_failed (GCP token mint failure → 502), or the management-supplied code on llm_limit_check. These surface verbatim; arbitrary middleware text never reaches the wire.

Things to scrutinise

Correctness. llm_router model match treats an empty Models slice as "claim every model" (middleware.go:238–248) for gateway-style providers — confirm no real provider record ships with an empty Models by accident. Path-prefix tie-break falls back to declaration order when no candidate prefix-matches, so the synthesiser must emit a deterministic order. llm_limit_record discards strconv.ParseInt errors (middleware.go:78–80) — relies on llm_response_parser always emitting parseable values; spot-check the streaming partial path on truncated bodies.

Security. Auth headers must NEVER appear on Mutations.HeadersAdd/Remove for the router — a direct headers path would bypass the framework gate. The capture-pointer handling is the kind of place a bug ships PII to logs silently; every synthesiser config path must set the pointer explicitly. llm_identity_inject body inject silently skips on a non-object metadata field (middleware.go:262–270) — header path still attributes, but body-level tag-budget enforcement doesn't run for that request.

Concurrency. cost_meter's two pricing tables are built once from the middleware config and never mutated, so the lookup path needs no lock or atomic swap — a price change replaces the whole instance. Every middleware is otherwise a stateless value receiver. Integration test uses real bufconn gRPC — race detector is the meaningful bar.

Perf. Hot path is lookupKV linear scan over <10 KVs; cost_meter.Cost is O(1); SSE accumulation is single-pass. No map allocation per call.

Observability. Every deny stamps llm_policy.decision=deny and a matching llm_policy.reason — access-log can pivot on either. llm_limit_record only logs at Debugf on RPC failure (middleware.go:125–130); operators need an alternate signal (metric on RecordLLMUsage failures) for counter accuracy.

Test coverage

FileTestsNotes
all_test.go1Registry surface lock
agentnetwork_chain_integration_test.go3Allow/deny/cap-exhaust vs live sqlite + bufconn gRPC
llm_request_parser/middleware_test.go18provider_id bypass, redaction, capture-pointer, rune-safe truncation
llm_router/middleware_test.go19Three-pass match, deny codes, path-prefix tie-break, header strip+inject
llm_limit_check/middleware_test.go6Allow/deny, fail-open on nil mgmt / RPC error, attribution stamping
llm_identity_inject/middleware_test.go28HeaderPair, JSONMetadata, ExtraHeaders, body inject, anti-spoof
llm_guardrail/middleware_test.go15Allowlist case-insensitivity, prompt capture toggle, deny shape
llm_guardrail/redact_test.go15Email, SSN, phone (E.164 + NA), bearer, IPv4; fixture-driven
llm_response_parser/middleware_test.go18Buffered OAI+Anthro, capture-pointer, redact, truncation
llm_response_parser/streaming_test.go7OAI usage frame, Anthro message_delta, truncated body best-effort
cost_meter/middleware_test.go22Each skip reason, provider-shape formulas, config-delivered defaults, per-record-beats-defaults + miss-falls-back, per-record uses surface formula, nil-pricing skips everything, invalid-rate rejection
llm_limit_record/middleware_test.go7Skip-on-no-signal, skip-on-missing-attribution, RPC failure swallowed

Cross-references

  • Sibling: 32-proxy-llm-parsers.md — SDK adapters
    • SSE framer + pricing table and cost formula.
  • Path-routed providers (Vertex AI + Bedrock), keyfile:: credential, GCP token minting, /bedrock prefix: 50-path-routed-providers.md.
  • Upstream config: management/server/agentnetwork/synthesizer (out of scope).
  • Framework: proxy/internal/middleware/{chain,dispatcher,accumulator,registry}.go.
  • Metadata key registry: proxy/internal/middleware/keys.go.
  • gRPC surface: proto.ProxyServiceClient.{CheckLLMPolicyLimits,RecordLLMUsage}.