mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-16 19:59:07 +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,55 @@
|
||||
package llm_request_parser
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestNormalizeBedrockModel(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"eu.anthropic.claude-sonnet-4-5-20250929-v1:0": "anthropic.claude-sonnet-4-5",
|
||||
"us.anthropic.claude-opus-4-8-20250101-v1:0": "anthropic.claude-opus-4-8",
|
||||
"apac.anthropic.claude-haiku-4-5-v1:0": "anthropic.claude-haiku-4-5",
|
||||
"anthropic.claude-sonnet-4-5-20250929-v1:0": "anthropic.claude-sonnet-4-5",
|
||||
"meta.llama3-3-70b-instruct-v1:0": "meta.llama3-3-70b-instruct",
|
||||
"amazon.nova-pro-v1:0": "amazon.nova-pro",
|
||||
"amazon.nova-2-lite-v1:0": "amazon.nova-2-lite",
|
||||
// Inference-profile ARN — model id lives in the last path segment.
|
||||
"arn:aws:bedrock:eu-central-1:123456789012:inference-profile/eu.anthropic.claude-sonnet-4-5-20250929-v1:0": "anthropic.claude-sonnet-4-5",
|
||||
}
|
||||
for in, want := range cases {
|
||||
require.Equal(t, want, normalizeBedrockModel(in), "normalize %q", in)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseBedrockPath(t *testing.T) {
|
||||
tests := []struct {
|
||||
path string
|
||||
model string
|
||||
stream bool
|
||||
ok bool
|
||||
}{
|
||||
{"/model/eu.anthropic.claude-sonnet-4-5-20250929-v1:0/invoke", "anthropic.claude-sonnet-4-5", false, true},
|
||||
{"/model/eu.anthropic.claude-sonnet-4-5-20250929-v1:0/invoke-with-response-stream", "anthropic.claude-sonnet-4-5", true, true},
|
||||
{"/model/eu.anthropic.claude-sonnet-4-5-20250929-v1:0/converse", "anthropic.claude-sonnet-4-5", false, true},
|
||||
{"/model/eu.anthropic.claude-sonnet-4-5-20250929-v1:0/converse-stream", "anthropic.claude-sonnet-4-5", true, true},
|
||||
// URL-encoded colon in the version suffix.
|
||||
{"/model/eu.anthropic.claude-sonnet-4-5-20250929-v1%3A0/invoke", "anthropic.claude-sonnet-4-5", false, true},
|
||||
// Optional "/bedrock" gateway-namespace prefix.
|
||||
{"/bedrock/model/eu.anthropic.claude-sonnet-4-5-20250929-v1:0/invoke-with-response-stream", "anthropic.claude-sonnet-4-5", true, true},
|
||||
{"/bedrock/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse", "anthropic.claude-sonnet-4-5", false, true},
|
||||
{"/v1/chat/completions", "", false, false},
|
||||
{"/model/foo", "", false, false},
|
||||
{"/model//invoke", "", false, false},
|
||||
{"/model/x/unknown-action", "", false, false},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
br, ok := parseBedrockPath(tt.path)
|
||||
require.Equal(t, tt.ok, ok, "ok for %q", tt.path)
|
||||
if tt.ok {
|
||||
require.Equal(t, tt.model, br.model, "model for %q", tt.path)
|
||||
require.Equal(t, tt.stream, br.stream, "stream for %q", tt.path)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package llm_request_parser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/netbirdio/netbird/proxy/internal/middleware"
|
||||
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
|
||||
)
|
||||
|
||||
// config is the on-wire config envelope for the middleware.
|
||||
//
|
||||
// ProviderID, when set, names the parser to use directly (matched
|
||||
// against llm.ParserByName, e.g. "openai", "anthropic"). The
|
||||
// agent-network synthesiser stamps this so requests routed through a
|
||||
// synthesised provider service don't depend on URL-shape sniffing,
|
||||
// which is the only signal the middleware otherwise has.
|
||||
type config struct {
|
||||
ProviderID string `json:"provider_id,omitempty"`
|
||||
// RedactPii, when true, runs PII redaction over the captured raw prompt
|
||||
// before it is emitted as llm.request_prompt_raw — so the
|
||||
// agent-network access-log row does NOT carry raw emails / SSNs /
|
||||
// phone numbers even though the framework's per-key redactor (Scan)
|
||||
// doesn't cover those prompt-shaped patterns. Sourced by the
|
||||
// synthesiser from the account's redact_pii toggle.
|
||||
RedactPii bool `json:"redact_pii,omitempty"`
|
||||
// CapturePrompt gates emission of llm.request_prompt_raw. A nil pointer
|
||||
// preserves the legacy default (emit), so callers that don't know about
|
||||
// the toggle (or pre-existing tests with empty config) keep working.
|
||||
// The synthesiser sets this explicitly to the account's
|
||||
// enable_prompt_collection toggle: false here suppresses the key
|
||||
// entirely so the access-log row carries no prompt content at all,
|
||||
// independent of redact_pii (which only controls the form of the
|
||||
// content when it IS emitted).
|
||||
CapturePrompt *bool `json:"capture_prompt,omitempty"`
|
||||
}
|
||||
|
||||
// Factory builds llm_request_parser instances from raw config bytes.
|
||||
type Factory struct{}
|
||||
|
||||
// ID returns the registry identifier.
|
||||
func (Factory) ID() string { return ID }
|
||||
|
||||
// New constructs a middleware instance. Empty, null, and {} configs are
|
||||
// accepted; non-empty rawConfig that fails to unmarshal is rejected so
|
||||
// misconfigurations surface at chain build time.
|
||||
func (Factory) New(rawConfig []byte) (middleware.Middleware, error) {
|
||||
var cfg config
|
||||
if len(bytes.TrimSpace(rawConfig)) > 0 {
|
||||
// Strict decode: a typo'd field (e.g. "capture_prompts") must fail
|
||||
// chain build rather than silently fall back to the emit-everything
|
||||
// default and leak prompts.
|
||||
dec := json.NewDecoder(bytes.NewReader(rawConfig))
|
||||
dec.DisallowUnknownFields()
|
||||
if err := dec.Decode(&cfg); err != nil {
|
||||
return nil, fmt.Errorf("decode config: %w", err)
|
||||
}
|
||||
}
|
||||
// Default capturePrompt to true (legacy emission) when the field is
|
||||
// absent so non-agent-network callers and pre-toggle tests keep working.
|
||||
capturePrompt := true
|
||||
if cfg.CapturePrompt != nil {
|
||||
capturePrompt = *cfg.CapturePrompt
|
||||
}
|
||||
return middlewareImpl{providerID: cfg.ProviderID, redactPii: cfg.RedactPii, capturePrompt: capturePrompt}, nil
|
||||
}
|
||||
|
||||
func init() {
|
||||
builtin.Register(Factory{})
|
||||
}
|
||||
@@ -0,0 +1,453 @@
|
||||
// Package llm_request_parser implements the SlotOnRequest middleware
|
||||
// that detects the LLM provider from the request URL, parses the JSON
|
||||
// request body for model and streaming flags, and extracts the user
|
||||
// prompt text. Emitted metadata feeds downstream middlewares (guardrail,
|
||||
// cost meter) and the access-log terminal sink.
|
||||
package llm_request_parser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/url"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/netbirdio/netbird/proxy/internal/llm"
|
||||
"github.com/netbirdio/netbird/proxy/internal/middleware"
|
||||
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
|
||||
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_guardrail"
|
||||
)
|
||||
|
||||
// ID is the registry key for this middleware.
|
||||
const ID = "llm_request_parser"
|
||||
|
||||
// Version is reported via Middleware.Version().
|
||||
const Version = "1.0.0"
|
||||
|
||||
// maxPromptBytes caps llm.request_prompt_raw at a size that fits within
|
||||
// MaxMetadataValueBytes with headroom. Truncation is rune-safe.
|
||||
const maxPromptBytes = 3500
|
||||
|
||||
// middlewareImpl is the concrete implementation. providerID, when set,
|
||||
// names the parser to use directly (bypasses URL sniffing). It is empty
|
||||
// for non-agent-network targets, which fall back to DetectParser on the
|
||||
// request path.
|
||||
type middlewareImpl struct {
|
||||
providerID string
|
||||
redactPii bool
|
||||
capturePrompt bool
|
||||
}
|
||||
|
||||
// ID returns the registry identifier.
|
||||
func (middlewareImpl) ID() string { return ID }
|
||||
|
||||
// Version returns the implementation version.
|
||||
func (middlewareImpl) Version() string { return Version }
|
||||
|
||||
// Slot reports the request slot.
|
||||
func (middlewareImpl) Slot() middleware.Slot { return middleware.SlotOnRequest }
|
||||
|
||||
// AcceptedContentTypes restricts body inspection to JSON.
|
||||
func (middlewareImpl) AcceptedContentTypes() []string {
|
||||
return []string{"application/json"}
|
||||
}
|
||||
|
||||
// MetadataKeys lists the closed allowlist of keys this middleware emits.
|
||||
func (middlewareImpl) MetadataKeys() []string {
|
||||
return []string{
|
||||
middleware.KeyLLMProvider,
|
||||
middleware.KeyLLMModel,
|
||||
middleware.KeyLLMStream,
|
||||
middleware.KeyLLMRequestPromptRaw,
|
||||
middleware.KeyLLMCaptureTruncated,
|
||||
middleware.KeyLLMSessionID,
|
||||
}
|
||||
}
|
||||
|
||||
// MutationsSupported reports that this middleware never mutates.
|
||||
func (middlewareImpl) MutationsSupported() bool { return false }
|
||||
|
||||
// Close is a no-op; the middleware is stateless.
|
||||
func (middlewareImpl) Close() error { return nil }
|
||||
|
||||
// Invoke detects the LLM provider, parses request facts, and emits
|
||||
// metadata. Always returns DecisionAllow; never errors. Provider
|
||||
// selection prefers the configured providerID (synthesiser-stamped on
|
||||
// agent-network targets) so requests routed to a custom upstream URL
|
||||
// still resolve. Falls back to URL sniffing when no providerID is set.
|
||||
func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
|
||||
out := &middleware.Output{Decision: middleware.DecisionAllow}
|
||||
if in == nil {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Google Vertex AI carries the model + publisher (vendor) in the URL path,
|
||||
// not the body, so it needs a dedicated extraction path.
|
||||
if vx, okv := parseVertexPath(extractPath(in.URL)); okv {
|
||||
return m.invokeVertex(in, vx), nil
|
||||
}
|
||||
|
||||
// AWS Bedrock likewise carries the model in the URL path (/model/{id}/{action}).
|
||||
if br, okb := parseBedrockPath(extractPath(in.URL)); okb {
|
||||
return m.invokeBedrock(in, br), nil
|
||||
}
|
||||
|
||||
parser, ok := llm.ParserByName(m.providerID)
|
||||
if !ok {
|
||||
parser, ok = llm.DetectParser(extractPath(in.URL))
|
||||
}
|
||||
if !ok {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
md := []middleware.KV{
|
||||
{Key: middleware.KeyLLMProvider, Value: parser.ProviderName()},
|
||||
}
|
||||
|
||||
// Session id is an opaque grouping identifier, not prompt content, so
|
||||
// it's emitted regardless of the prompt-collection toggle — session
|
||||
// grouping must work even when prompt capture is off. Prefer a header
|
||||
// (Codex sends the session as an HTTP header, and headers survive an
|
||||
// oversized request whose body capture was bypassed) and resolve it
|
||||
// before ParseRequest so a malformed body still keeps the header id.
|
||||
sessionID := sessionIDFromHeaders(in.Headers)
|
||||
if sessionID == "" {
|
||||
sessionID = parser.ExtractSessionID(in.Body)
|
||||
}
|
||||
appendSessionID := func(md []middleware.KV) []middleware.KV {
|
||||
if sessionID != "" {
|
||||
return append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID})
|
||||
}
|
||||
return md
|
||||
}
|
||||
|
||||
facts, err := parser.ParseRequest(in.Body)
|
||||
if err != nil {
|
||||
if logger := builtin.Context().Logger; logger != nil {
|
||||
logger.Debugf("llm_request_parser: parse request body: %v", err)
|
||||
}
|
||||
md = appendSessionID(md)
|
||||
md = appendCaptureTruncated(md, false, in.BodyTruncated)
|
||||
out.Metadata = md
|
||||
return out, nil
|
||||
}
|
||||
|
||||
if facts.Model != "" {
|
||||
md = append(md, middleware.KV{Key: middleware.KeyLLMModel, Value: facts.Model})
|
||||
}
|
||||
md = append(md, middleware.KV{Key: middleware.KeyLLMStream, Value: strconv.FormatBool(facts.Stream)})
|
||||
md = appendSessionID(md)
|
||||
|
||||
prompt, promptTruncated := truncatePrompt(parser.ExtractPrompt(in.Body))
|
||||
if prompt != "" && m.capturePrompt {
|
||||
if m.redactPii {
|
||||
// Apply redaction BEFORE the value lands in the metadata bag, so
|
||||
// the access-log row never carries raw emails / SSNs / phones.
|
||||
// The downstream llm_guardrail middleware reads this key to
|
||||
// produce llm.request_prompt; RedactPII is idempotent so its
|
||||
// second pass is a no-op. Redaction can grow the text, so
|
||||
// re-truncate to keep the value within the metadata cap.
|
||||
prompt = llm_guardrail.RedactPII(prompt)
|
||||
var redactedTruncated bool
|
||||
prompt, redactedTruncated = truncatePrompt(prompt)
|
||||
promptTruncated = promptTruncated || redactedTruncated
|
||||
}
|
||||
md = append(md, middleware.KV{Key: middleware.KeyLLMRequestPromptRaw, Value: prompt})
|
||||
}
|
||||
|
||||
md = appendCaptureTruncated(md, promptTruncated, in.BodyTruncated)
|
||||
out.Metadata = md
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// sessionIDHeaders are request header names that may carry a client
|
||||
// session identifier, checked in order, case-insensitively. Matching is
|
||||
// against Go's canonical header form, so use the hyphenated names the
|
||||
// clients actually send: "x-claude-code-session-id" (Claude Code),
|
||||
// "session-id" (OpenAI Codex — confirmed on the wire as "Session-Id"),
|
||||
// and "x-session-id" as a generic convention.
|
||||
var sessionIDHeaders = []string{"x-claude-code-session-id", "session-id", "x-session-id"}
|
||||
|
||||
// sessionIDFromHeaders returns the first non-empty value among the known
|
||||
// session header names, or "" when none is present. Headers arrive in
|
||||
// canonical form, so the match is case-insensitive.
|
||||
func sessionIDFromHeaders(headers []middleware.KV) string {
|
||||
for _, want := range sessionIDHeaders {
|
||||
for _, kv := range headers {
|
||||
if strings.EqualFold(kv.Key, want) && kv.Value != "" {
|
||||
return kv.Value
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// appendCaptureTruncated stamps the capture_truncated marker reflecting
|
||||
// either prompt-side truncation or upstream body truncation.
|
||||
func appendCaptureTruncated(md []middleware.KV, promptTruncated, bodyTruncated bool) []middleware.KV {
|
||||
value := "false"
|
||||
if promptTruncated || bodyTruncated {
|
||||
value = "true"
|
||||
}
|
||||
return append(md, middleware.KV{Key: middleware.KeyLLMCaptureTruncated, Value: value})
|
||||
}
|
||||
|
||||
// truncatePrompt clamps a prompt string to maxPromptBytes on a UTF-8
|
||||
// rune boundary. Returns the clamped string and whether truncation
|
||||
// occurred.
|
||||
func truncatePrompt(s string) (string, bool) {
|
||||
if len(s) <= maxPromptBytes {
|
||||
return s, false
|
||||
}
|
||||
cut := maxPromptBytes
|
||||
for cut > 0 && !utf8.RuneStart(s[cut]) {
|
||||
cut--
|
||||
}
|
||||
return s[:cut], true
|
||||
}
|
||||
|
||||
// extractPath returns the path component of a URL that may be absolute
|
||||
// or already a path. Parse errors fall back to the raw input.
|
||||
func extractPath(raw string) string {
|
||||
if raw == "" {
|
||||
return ""
|
||||
}
|
||||
u, err := url.Parse(raw)
|
||||
if err != nil || u.Path == "" {
|
||||
return raw
|
||||
}
|
||||
return u.Path
|
||||
}
|
||||
|
||||
// vertexRequest is the model + vendor extracted from a Vertex AI publisher
|
||||
// path (the model is in the URL, not the body).
|
||||
type vertexRequest struct {
|
||||
publisher string
|
||||
model string
|
||||
stream bool
|
||||
}
|
||||
|
||||
// parseVertexPath extracts the publisher, model, and streaming flag from a
|
||||
// Vertex publisher endpoint:
|
||||
//
|
||||
// /v1/projects/{project}/locations/{region}/publishers/{publisher}/models/{model}:{action}
|
||||
//
|
||||
// The model's "@version" suffix is stripped so it matches catalog/pricing.
|
||||
func parseVertexPath(reqPath string) (vertexRequest, bool) {
|
||||
const pubSep, modSep = "/publishers/", "/models/"
|
||||
if !strings.HasPrefix(reqPath, "/v1/projects/") {
|
||||
return vertexRequest{}, false
|
||||
}
|
||||
pubIdx := strings.Index(reqPath, pubSep)
|
||||
modIdx := strings.Index(reqPath, modSep)
|
||||
if pubIdx < 0 || modIdx <= pubIdx {
|
||||
return vertexRequest{}, false
|
||||
}
|
||||
publisher := reqPath[pubIdx+len(pubSep) : modIdx]
|
||||
rest := reqPath[modIdx+len(modSep):] // {model}:{action}
|
||||
if publisher == "" || rest == "" {
|
||||
return vertexRequest{}, false
|
||||
}
|
||||
model, action := rest, ""
|
||||
if c := strings.LastIndex(rest, ":"); c >= 0 {
|
||||
model, action = rest[:c], rest[c+1:]
|
||||
}
|
||||
if at := strings.Index(model, "@"); at >= 0 {
|
||||
model = model[:at]
|
||||
}
|
||||
if model == "" {
|
||||
return vertexRequest{}, false
|
||||
}
|
||||
return vertexRequest{publisher: publisher, model: model, stream: strings.HasPrefix(action, "stream")}, true
|
||||
}
|
||||
|
||||
// vertexPublisherVendor maps a Vertex publisher to the parser surface its
|
||||
// requests/responses speak. Empty for publishers without a parser yet
|
||||
// (e.g. google/gemini) — the request still routes, but isn't metered.
|
||||
func vertexPublisherVendor(publisher string) string {
|
||||
switch strings.ToLower(publisher) {
|
||||
case "anthropic":
|
||||
return "anthropic"
|
||||
case "openai":
|
||||
return "openai"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// invokeVertex emits the model/vendor/session/prompt for a Vertex publisher
|
||||
// request, using the publisher's parser to read the (vendor-native) body.
|
||||
func (m middlewareImpl) invokeVertex(in *middleware.Input, vx vertexRequest) *middleware.Output {
|
||||
out := &middleware.Output{Decision: middleware.DecisionAllow}
|
||||
vendor := vertexPublisherVendor(vx.publisher)
|
||||
|
||||
md := []middleware.KV{}
|
||||
if vendor != "" {
|
||||
md = append(md, middleware.KV{Key: middleware.KeyLLMProvider, Value: vendor})
|
||||
}
|
||||
md = append(md, middleware.KV{Key: middleware.KeyLLMModel, Value: vx.model})
|
||||
md = append(md, middleware.KV{Key: middleware.KeyLLMStream, Value: strconv.FormatBool(vx.stream)})
|
||||
|
||||
var parser llm.Parser
|
||||
if vendor != "" {
|
||||
parser, _ = llm.ParserByName(vendor)
|
||||
}
|
||||
|
||||
sessionID := sessionIDFromHeaders(in.Headers)
|
||||
if sessionID == "" && parser != nil {
|
||||
sessionID = parser.ExtractSessionID(in.Body)
|
||||
}
|
||||
if sessionID != "" {
|
||||
md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID})
|
||||
}
|
||||
|
||||
promptTruncated := false
|
||||
if parser != nil && m.capturePrompt {
|
||||
var prompt string
|
||||
prompt, promptTruncated = truncatePrompt(parser.ExtractPrompt(in.Body))
|
||||
if prompt != "" {
|
||||
if m.redactPii {
|
||||
prompt = llm_guardrail.RedactPII(prompt)
|
||||
var rt bool
|
||||
prompt, rt = truncatePrompt(prompt)
|
||||
promptTruncated = promptTruncated || rt
|
||||
}
|
||||
md = append(md, middleware.KV{Key: middleware.KeyLLMRequestPromptRaw, Value: prompt})
|
||||
}
|
||||
}
|
||||
md = appendCaptureTruncated(md, promptTruncated, in.BodyTruncated)
|
||||
out.Metadata = md
|
||||
return out
|
||||
}
|
||||
|
||||
// bedrockRequest is the model + streaming flag extracted from an AWS Bedrock
|
||||
// model path. The InvokeModel vs Converse distinction is recovered downstream
|
||||
// from the response body shape, so only the streaming flag is carried here.
|
||||
type bedrockRequest struct {
|
||||
model string
|
||||
stream bool
|
||||
}
|
||||
|
||||
// bedrockNamespacePrefix is an optional gateway-namespace prefix some clients
|
||||
// put before the native Bedrock path to disambiguate it from other providers
|
||||
// that also use "/model/...".
|
||||
const bedrockNamespacePrefix = "/bedrock"
|
||||
|
||||
// trimBedrockNamespace removes an optional "/bedrock" namespace prefix, leaving
|
||||
// the native Bedrock path ("/model/...").
|
||||
func trimBedrockNamespace(reqPath string) string {
|
||||
if strings.HasPrefix(reqPath, bedrockNamespacePrefix+"/") {
|
||||
return strings.TrimPrefix(reqPath, bedrockNamespacePrefix)
|
||||
}
|
||||
return reqPath
|
||||
}
|
||||
|
||||
// bedrockRegionPrefixes are the cross-region inference-profile prefixes that
|
||||
// front a Bedrock model id (e.g. "eu.anthropic.claude-...").
|
||||
var bedrockRegionPrefixes = []string{"us.", "eu.", "apac.", "global."}
|
||||
|
||||
// bedrockVersionSuffix matches the trailing "-vN[:N]" or "-YYYYMMDD-vN[:N]"
|
||||
// version/throughput suffix of a Bedrock model id.
|
||||
var bedrockVersionSuffix = regexp.MustCompile(`-(\d{8}-)?v\d+(:\d+)?$`)
|
||||
|
||||
// parseBedrockPath extracts the model and streaming/converse flags from an AWS
|
||||
// Bedrock runtime model endpoint:
|
||||
//
|
||||
// /model/{modelId}/{action}
|
||||
//
|
||||
// action ∈ {invoke, invoke-with-response-stream, converse, converse-stream}.
|
||||
// The modelId may be URL-encoded and may carry a cross-region inference-profile
|
||||
// prefix and a version suffix; normalizeBedrockModel strips both so the model
|
||||
// matches catalog pricing.
|
||||
func parseBedrockPath(reqPath string) (bedrockRequest, bool) {
|
||||
reqPath = trimBedrockNamespace(reqPath)
|
||||
const prefix = "/model/"
|
||||
if !strings.HasPrefix(reqPath, prefix) {
|
||||
return bedrockRequest{}, false
|
||||
}
|
||||
rest := reqPath[len(prefix):]
|
||||
slash := strings.LastIndex(rest, "/")
|
||||
if slash <= 0 || slash == len(rest)-1 {
|
||||
return bedrockRequest{}, false
|
||||
}
|
||||
rawModel, action := rest[:slash], rest[slash+1:]
|
||||
if decoded, err := url.PathUnescape(rawModel); err == nil {
|
||||
rawModel = decoded
|
||||
}
|
||||
model := normalizeBedrockModel(rawModel)
|
||||
if model == "" {
|
||||
return bedrockRequest{}, false
|
||||
}
|
||||
switch action {
|
||||
case "invoke", "converse":
|
||||
return bedrockRequest{model: model}, true
|
||||
case "invoke-with-response-stream", "converse-stream":
|
||||
return bedrockRequest{model: model, stream: true}, true
|
||||
default:
|
||||
return bedrockRequest{}, false
|
||||
}
|
||||
}
|
||||
|
||||
// normalizeBedrockModel strips an ARN wrapper, a cross-region inference-profile
|
||||
// prefix, and the version/throughput suffix from a Bedrock model id so it
|
||||
// matches the catalog/pricing key, e.g.
|
||||
// "eu.anthropic.claude-sonnet-4-5-20250929-v1:0" -> "anthropic.claude-sonnet-4-5"
|
||||
// and "arn:aws:bedrock:eu-central-1:123:inference-profile/eu.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
// -> "anthropic.claude-sonnet-4-5".
|
||||
func normalizeBedrockModel(modelID string) string {
|
||||
m := modelID
|
||||
// A full ARN (inference-profile / provisioned-throughput / foundation-model)
|
||||
// carries the model id in its last path segment.
|
||||
if strings.HasPrefix(m, "arn:") {
|
||||
if i := strings.LastIndex(m, "/"); i >= 0 {
|
||||
m = m[i+1:]
|
||||
}
|
||||
}
|
||||
for _, p := range bedrockRegionPrefixes {
|
||||
if strings.HasPrefix(m, p) {
|
||||
m = m[len(p):]
|
||||
break
|
||||
}
|
||||
}
|
||||
return bedrockVersionSuffix.ReplaceAllString(m, "")
|
||||
}
|
||||
|
||||
// invokeBedrock emits the model/provider/session/prompt for an AWS Bedrock
|
||||
// request. Bedrock is metered under the dedicated "bedrock" parser, which reads
|
||||
// both the InvokeModel and Converse response shapes.
|
||||
func (m middlewareImpl) invokeBedrock(in *middleware.Input, br bedrockRequest) *middleware.Output {
|
||||
out := &middleware.Output{Decision: middleware.DecisionAllow}
|
||||
md := []middleware.KV{
|
||||
{Key: middleware.KeyLLMProvider, Value: llm.ProviderNameBedrock},
|
||||
{Key: middleware.KeyLLMModel, Value: br.model},
|
||||
{Key: middleware.KeyLLMStream, Value: strconv.FormatBool(br.stream)},
|
||||
}
|
||||
|
||||
parser, _ := llm.ParserByName(llm.ProviderNameBedrock)
|
||||
sessionID := sessionIDFromHeaders(in.Headers)
|
||||
if sessionID == "" && parser != nil {
|
||||
sessionID = parser.ExtractSessionID(in.Body)
|
||||
}
|
||||
if sessionID != "" {
|
||||
md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID})
|
||||
}
|
||||
|
||||
promptTruncated := false
|
||||
if parser != nil && m.capturePrompt {
|
||||
var prompt string
|
||||
prompt, promptTruncated = truncatePrompt(parser.ExtractPrompt(in.Body))
|
||||
if prompt != "" {
|
||||
if m.redactPii {
|
||||
prompt = llm_guardrail.RedactPII(prompt)
|
||||
var rt bool
|
||||
prompt, rt = truncatePrompt(prompt)
|
||||
promptTruncated = promptTruncated || rt
|
||||
}
|
||||
md = append(md, middleware.KV{Key: middleware.KeyLLMRequestPromptRaw, Value: prompt})
|
||||
}
|
||||
}
|
||||
md = appendCaptureTruncated(md, promptTruncated, in.BodyTruncated)
|
||||
out.Metadata = md
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,418 @@
|
||||
package llm_request_parser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/proxy/internal/middleware"
|
||||
)
|
||||
|
||||
func metaValue(t *testing.T, kvs []middleware.KV, key string) (string, bool) {
|
||||
t.Helper()
|
||||
for _, kv := range kvs {
|
||||
if kv.Key == key {
|
||||
return kv.Value, true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func newMiddleware(t *testing.T) middleware.Middleware {
|
||||
t.Helper()
|
||||
mw, err := Factory{}.New(nil)
|
||||
require.NoError(t, err, "factory must accept nil config")
|
||||
return mw
|
||||
}
|
||||
|
||||
func TestMiddleware_StaticSurface(t *testing.T) {
|
||||
mw := newMiddleware(t)
|
||||
assert.Equal(t, ID, mw.ID(), "ID must match the registered constant")
|
||||
assert.Equal(t, Version, mw.Version(), "Version must match the constant")
|
||||
assert.Equal(t, middleware.SlotOnRequest, mw.Slot(), "must run in the request slot")
|
||||
assert.Equal(t, []string{"application/json"}, mw.AcceptedContentTypes(), "only JSON bodies are needed")
|
||||
assert.False(t, mw.MutationsSupported(), "request parser never mutates")
|
||||
assert.NoError(t, mw.Close(), "Close on stateless middleware is a no-op")
|
||||
|
||||
keys := mw.MetadataKeys()
|
||||
expected := []string{
|
||||
middleware.KeyLLMProvider,
|
||||
middleware.KeyLLMModel,
|
||||
middleware.KeyLLMStream,
|
||||
middleware.KeyLLMRequestPromptRaw,
|
||||
middleware.KeyLLMCaptureTruncated,
|
||||
middleware.KeyLLMSessionID,
|
||||
}
|
||||
assert.Equal(t, expected, keys, "metadata key allowlist must match the spec")
|
||||
}
|
||||
|
||||
func TestFactory_AcceptsEmptyAndJSONConfig(t *testing.T) {
|
||||
cases := [][]byte{nil, {}, []byte("null"), []byte("{}"), []byte(" ")}
|
||||
for _, raw := range cases {
|
||||
mw, err := Factory{}.New(raw)
|
||||
require.NoError(t, err, "empty/null/object config must be accepted")
|
||||
require.NotNil(t, mw, "factory must return a middleware instance")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFactory_RejectsMalformedConfig(t *testing.T) {
|
||||
mw, err := Factory{}.New([]byte("{not json"))
|
||||
require.Error(t, err, "malformed config must surface at construction")
|
||||
assert.Nil(t, mw, "no instance is returned on error")
|
||||
}
|
||||
|
||||
func TestInvoke_OpenAIBufferedChatCompletion(t *testing.T) {
|
||||
mw := newMiddleware(t)
|
||||
body := []byte(`{"model":"gpt-4o-mini","stream":false,"messages":[{"role":"user","content":"Hello, world!"}]}`)
|
||||
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/chat/completions",
|
||||
Body: body,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out, "output must be returned")
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "request parser always allows")
|
||||
|
||||
provider, ok := metaValue(t, out.Metadata, middleware.KeyLLMProvider)
|
||||
require.True(t, ok, "provider metadata must be set")
|
||||
assert.Equal(t, "openai", provider, "OpenAI provider detected from path")
|
||||
|
||||
model, ok := metaValue(t, out.Metadata, middleware.KeyLLMModel)
|
||||
require.True(t, ok, "model metadata must be set")
|
||||
assert.Equal(t, "gpt-4o-mini", model, "model echoed from request body")
|
||||
|
||||
stream, ok := metaValue(t, out.Metadata, middleware.KeyLLMStream)
|
||||
require.True(t, ok, "stream metadata must be set")
|
||||
assert.Equal(t, "false", stream, "buffered request reports stream=false")
|
||||
|
||||
prompt, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPromptRaw)
|
||||
require.True(t, ok, "prompt metadata must be set when extractable")
|
||||
assert.Contains(t, prompt, "Hello, world!", "extracted prompt carries the user message")
|
||||
|
||||
truncated, ok := metaValue(t, out.Metadata, middleware.KeyLLMCaptureTruncated)
|
||||
require.True(t, ok, "capture_truncated must always be emitted on success")
|
||||
assert.Equal(t, "false", truncated, "no truncation on a small body")
|
||||
}
|
||||
|
||||
func TestInvoke_EmitsSessionID(t *testing.T) {
|
||||
mw := newMiddleware(t)
|
||||
|
||||
t.Run("codex session from client_metadata", func(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-5.5","client_metadata":{"session_id":"sess-codex-1"},"input":[]}`)
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{URL: "/v1/responses", Body: body})
|
||||
require.NoError(t, err)
|
||||
sid, ok := metaValue(t, out.Metadata, middleware.KeyLLMSessionID)
|
||||
require.True(t, ok, "session id must be emitted for Codex requests")
|
||||
assert.Equal(t, "sess-codex-1", sid, "session id must come from client_metadata.session_id")
|
||||
})
|
||||
|
||||
t.Run("no session id key when absent", func(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{URL: "/v1/chat/completions", Body: body})
|
||||
require.NoError(t, err)
|
||||
_, ok := metaValue(t, out.Metadata, middleware.KeyLLMSessionID)
|
||||
assert.False(t, ok, "no session id key emitted when the request carries none")
|
||||
})
|
||||
|
||||
t.Run("claude code session header", func(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-opus-4-8","messages":[{"role":"user","content":"hi"}]}`)
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/messages",
|
||||
Body: body,
|
||||
Headers: []middleware.KV{{Key: "X-Claude-Code-Session-Id", Value: "cc-sess-1"}},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
sid, ok := metaValue(t, out.Metadata, middleware.KeyLLMSessionID)
|
||||
require.True(t, ok, "Claude Code session id must be read from X-Claude-Code-Session-Id")
|
||||
assert.Equal(t, "cc-sess-1", sid, "session id must come from the Claude Code session header")
|
||||
})
|
||||
|
||||
t.Run("codex Session-Id header", func(t *testing.T) {
|
||||
// Codex sends the session as the canonical header "Session-Id".
|
||||
body := []byte(`{"model":"gpt-5.5","input":[]}`)
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/responses",
|
||||
Body: body,
|
||||
Headers: []middleware.KV{{Key: "Session-Id", Value: "sess-hdr-1"}},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
sid, ok := metaValue(t, out.Metadata, middleware.KeyLLMSessionID)
|
||||
require.True(t, ok, "session id must be read from the Session-Id header")
|
||||
assert.Equal(t, "sess-hdr-1", sid, "session id must come from the Codex Session-Id header")
|
||||
})
|
||||
|
||||
t.Run("header wins over body and survives bypassed body", func(t *testing.T) {
|
||||
// Oversized request: body was bypassed to a routing stub with no
|
||||
// client_metadata, but the header still carries the session.
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/responses",
|
||||
Body: []byte(`{"model":"gpt-5.5","stream":true}`),
|
||||
Headers: []middleware.KV{{Key: "X-Session-Id", Value: "sess-hdr-2"}},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
sid, _ := metaValue(t, out.Metadata, middleware.KeyLLMSessionID)
|
||||
assert.Equal(t, "sess-hdr-2", sid, "x-session-id header must be honoured when the body carries no marker")
|
||||
})
|
||||
}
|
||||
|
||||
func TestInvoke_OpenAIStreamingChatCompletion(t *testing.T) {
|
||||
mw := newMiddleware(t)
|
||||
body := []byte(`{"model":"gpt-4o-mini","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
||||
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/chat/completions",
|
||||
Body: body,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
stream, ok := metaValue(t, out.Metadata, middleware.KeyLLMStream)
|
||||
require.True(t, ok, "stream metadata must be set")
|
||||
assert.Equal(t, "true", stream, "stream flag echoed for SSE-bound request")
|
||||
}
|
||||
|
||||
func TestInvoke_AnthropicMessages(t *testing.T) {
|
||||
mw := newMiddleware(t)
|
||||
body := []byte(`{"model":"claude-sonnet-4-5","stream":false,"messages":[{"role":"user","content":"What is the weather?"}]}`)
|
||||
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/messages",
|
||||
Body: body,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
provider, ok := metaValue(t, out.Metadata, middleware.KeyLLMProvider)
|
||||
require.True(t, ok, "provider metadata must be set")
|
||||
assert.Equal(t, "anthropic", provider, "Anthropic provider detected from path")
|
||||
|
||||
model, _ := metaValue(t, out.Metadata, middleware.KeyLLMModel)
|
||||
assert.Equal(t, "claude-sonnet-4-5", model, "anthropic model echoed")
|
||||
|
||||
prompt, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPromptRaw)
|
||||
require.True(t, ok, "prompt metadata must be set")
|
||||
assert.Contains(t, prompt, "What is the weather?", "anthropic message text extracted")
|
||||
}
|
||||
|
||||
func TestInvoke_UnknownURLNoMetadata(t *testing.T) {
|
||||
mw := newMiddleware(t)
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/healthz",
|
||||
Body: []byte(`{"model":"x"}`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "unknown paths still allow")
|
||||
assert.Empty(t, out.Metadata, "no metadata is emitted when no parser matches")
|
||||
}
|
||||
|
||||
func TestInvoke_ProviderIDConfigBypassesURLSniff(t *testing.T) {
|
||||
mw, err := Factory{}.New([]byte(`{"provider_id":"openai"}`))
|
||||
require.NoError(t, err, "factory must accept provider_id config")
|
||||
|
||||
// URL doesn't match any of the OpenAI path hints — the provider_id
|
||||
// config is the only signal the middleware has.
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/custom/gateway/foo/bar",
|
||||
Body: []byte(`{"model":"gpt-4o-mini","messages":[{"role":"user","content":"Hi"}]}`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision)
|
||||
|
||||
provider, ok := metaValue(t, out.Metadata, middleware.KeyLLMProvider)
|
||||
require.True(t, ok, "provider must be emitted when provider_id is configured even on unknown URLs")
|
||||
assert.Equal(t, "openai", provider, "provider_id config selects the OpenAI parser")
|
||||
|
||||
model, ok := metaValue(t, out.Metadata, middleware.KeyLLMModel)
|
||||
require.True(t, ok, "model still extracted from the body")
|
||||
assert.Equal(t, "gpt-4o-mini", model)
|
||||
}
|
||||
|
||||
func TestInvoke_UnknownProviderIDFallsBackToURL(t *testing.T) {
|
||||
mw, err := Factory{}.New([]byte(`{"provider_id":"not-a-real-parser"}`))
|
||||
require.NoError(t, err, "factory must accept any provider_id string")
|
||||
|
||||
// URL hits the OpenAI surface, so URL sniffing should still resolve
|
||||
// even though the configured provider_id doesn't match a parser.
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/chat/completions",
|
||||
Body: []byte(`{"model":"gpt-4o-mini"}`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
|
||||
provider, ok := metaValue(t, out.Metadata, middleware.KeyLLMProvider)
|
||||
require.True(t, ok, "fallback URL sniffing must populate the provider")
|
||||
assert.Equal(t, "openai", provider)
|
||||
}
|
||||
|
||||
func TestInvoke_MalformedBodyAllowsWithProvider(t *testing.T) {
|
||||
mw := newMiddleware(t)
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/chat/completions",
|
||||
Body: []byte(`{not json`),
|
||||
})
|
||||
require.NoError(t, err, "malformed body must not error")
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "decision is always allow")
|
||||
|
||||
provider, ok := metaValue(t, out.Metadata, middleware.KeyLLMProvider)
|
||||
require.True(t, ok, "provider metadata is emitted before body parse")
|
||||
assert.Equal(t, "openai", provider, "provider stays even when body parse fails")
|
||||
|
||||
_, hasModel := metaValue(t, out.Metadata, middleware.KeyLLMModel)
|
||||
assert.False(t, hasModel, "no model metadata when parse fails")
|
||||
|
||||
truncated, ok := metaValue(t, out.Metadata, middleware.KeyLLMCaptureTruncated)
|
||||
require.True(t, ok, "capture_truncated is emitted on parse error path")
|
||||
assert.Equal(t, "false", truncated, "no truncation marker without truncated body or prompt")
|
||||
}
|
||||
|
||||
func TestInvoke_TruncatesLongPrompt(t *testing.T) {
|
||||
mw := newMiddleware(t)
|
||||
long := strings.Repeat("x", maxPromptBytes*2)
|
||||
body := []byte(`{"model":"gpt-4o-mini","messages":[{"role":"user","content":"` + long + `"}]}`)
|
||||
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/chat/completions",
|
||||
Body: body,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
prompt, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPromptRaw)
|
||||
require.True(t, ok, "prompt metadata must be set")
|
||||
assert.LessOrEqual(t, len(prompt), maxPromptBytes, "prompt must respect the byte budget")
|
||||
|
||||
truncated, ok := metaValue(t, out.Metadata, middleware.KeyLLMCaptureTruncated)
|
||||
require.True(t, ok, "capture_truncated must be set")
|
||||
assert.Equal(t, "true", truncated, "truncation marker raised when prompt is clipped")
|
||||
}
|
||||
|
||||
func TestInvoke_TruncatesOnRuneBoundary(t *testing.T) {
|
||||
mw := newMiddleware(t)
|
||||
// Each ☃ is 3 bytes in UTF-8; build a string whose byte length exceeds
|
||||
// maxPromptBytes with snowmen straddling the cut point.
|
||||
rune3 := "☃"
|
||||
repeats := (maxPromptBytes / len(rune3)) + 5
|
||||
long := strings.Repeat(rune3, repeats)
|
||||
body := []byte(`{"model":"gpt-4o-mini","messages":[{"role":"user","content":"` + long + `"}]}`)
|
||||
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/chat/completions",
|
||||
Body: body,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
prompt, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPromptRaw)
|
||||
require.True(t, ok, "prompt metadata must be set")
|
||||
assert.LessOrEqual(t, len(prompt), maxPromptBytes, "prompt must respect the byte budget")
|
||||
assert.True(t, strings.HasSuffix(prompt, rune3) || !strings.ContainsRune(prompt[len(prompt)-1:], 0xFFFD),
|
||||
"truncation must not split a multi-byte rune")
|
||||
}
|
||||
|
||||
func TestInvoke_BodyTruncatedRaisesCaptureTruncated(t *testing.T) {
|
||||
mw := newMiddleware(t)
|
||||
body := []byte(`{"model":"gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`)
|
||||
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{
|
||||
URL: "/v1/chat/completions",
|
||||
Body: body,
|
||||
BodyTruncated: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
truncated, ok := metaValue(t, out.Metadata, middleware.KeyLLMCaptureTruncated)
|
||||
require.True(t, ok, "capture_truncated must be set")
|
||||
assert.Equal(t, "true", truncated, "BodyTruncated input flips the marker even when prompt fits")
|
||||
}
|
||||
|
||||
// TestInvoke_RedactPii_RedactsBeforeEmittingRawPrompt covers the GC contract:
|
||||
// when the synthesiser sets redact_pii=true on the parser config, the value
|
||||
// emitted as llm.request_prompt_raw must already be redacted, so the
|
||||
// access-log row never carries raw emails / SSNs / phones — even though the
|
||||
// downstream llm_guardrail middleware also runs.
|
||||
func TestInvoke_RedactPii_RedactsBeforeEmittingRawPrompt(t *testing.T) {
|
||||
mw, err := Factory{}.New([]byte(`{"redact_pii":true}`))
|
||||
require.NoError(t, err)
|
||||
|
||||
body := []byte(`{"model":"gpt-4o-mini","stream":false,"messages":[{"role":"user","content":"contact alice.johnson@example.com SSN 123-45-6789 phone (202) 555-0147 and bob 202/555/0108"}]}`)
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{URL: "/v1/chat/completions", Body: body})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
|
||||
raw, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPromptRaw)
|
||||
require.True(t, ok, "raw prompt key must still be emitted")
|
||||
assert.Contains(t, raw, "[REDACTED:email]", "email must be redacted before emit")
|
||||
assert.Contains(t, raw, "[REDACTED:ssn]", "ssn must be redacted before emit")
|
||||
assert.Contains(t, raw, "[REDACTED:phone]", "phone must be redacted before emit")
|
||||
assert.NotContains(t, raw, "alice.johnson@example.com", "raw email must not survive")
|
||||
assert.NotContains(t, raw, "123-45-6789", "raw SSN must not survive")
|
||||
assert.NotContains(t, raw, "(202) 555-0147", "parenthesised phone must not survive")
|
||||
assert.NotContains(t, raw, "202/555/0108", "slash-separated phone must not survive")
|
||||
}
|
||||
|
||||
// TestInvoke_CapturePromptOff_DoesNotEmitRawPrompt covers the contract for
|
||||
// the account-level enable_prompt_collection toggle: when the synthesiser sets
|
||||
// capture_prompt=false (operator hasn't opted in to prompt content), the
|
||||
// parser MUST NOT emit llm.request_prompt_raw at all — otherwise the access
|
||||
// log carries the user's input even though log collection is meant to be
|
||||
// metadata-only (provider, model, tokens, cost). The other facts the parser
|
||||
// emits (provider, model, stream, capture_truncated) stay.
|
||||
func TestInvoke_CapturePromptOff_DoesNotEmitRawPrompt(t *testing.T) {
|
||||
mw, err := Factory{}.New([]byte(`{"capture_prompt":false}`))
|
||||
require.NoError(t, err)
|
||||
body := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"contact alice@example.com SSN 123-45-6789"}]}`)
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{URL: "/v1/chat/completions", Body: body})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
|
||||
_, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPromptRaw)
|
||||
assert.False(t, ok, "llm.request_prompt_raw must NOT be emitted when capture_prompt is false")
|
||||
// Non-content facts must still flow.
|
||||
_, ok = metaValue(t, out.Metadata, middleware.KeyLLMModel)
|
||||
assert.True(t, ok, "model fact must still be emitted")
|
||||
_, ok = metaValue(t, out.Metadata, middleware.KeyLLMProvider)
|
||||
assert.True(t, ok, "provider fact must still be emitted")
|
||||
}
|
||||
|
||||
// TestInvoke_CapturePromptUnset_PreservesLegacyEmission documents the default
|
||||
// behavior: an empty / legacy config (no capture_prompt field) keeps the
|
||||
// existing emission, so non-agent-network callers and pre-toggle tests don't
|
||||
// suddenly lose data.
|
||||
func TestInvoke_CapturePromptUnset_PreservesLegacyEmission(t *testing.T) {
|
||||
mw, err := Factory{}.New([]byte(`{}`))
|
||||
require.NoError(t, err)
|
||||
body := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hello"}]}`)
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{URL: "/v1/chat/completions", Body: body})
|
||||
require.NoError(t, err)
|
||||
_, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPromptRaw)
|
||||
assert.True(t, ok, "absent capture_prompt must preserve emission (backwards-compatible default)")
|
||||
}
|
||||
|
||||
// TestInvoke_RedactPii_OffShipsRawPrompt is the inverse: when redact_pii is
|
||||
// false (default) the operator opted out and the raw prompt is shipped
|
||||
// verbatim, so audit / debugging consumers still get the full body.
|
||||
func TestInvoke_RedactPii_OffShipsRawPrompt(t *testing.T) {
|
||||
mw, err := Factory{}.New([]byte(`{}`))
|
||||
require.NoError(t, err)
|
||||
|
||||
body := []byte(`{"model":"gpt-4o-mini","messages":[{"role":"user","content":"alice.johnson@example.com"}]}`)
|
||||
out, err := mw.Invoke(context.Background(), &middleware.Input{URL: "/v1/chat/completions", Body: body})
|
||||
require.NoError(t, err)
|
||||
|
||||
raw, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPromptRaw)
|
||||
require.True(t, ok)
|
||||
assert.Contains(t, raw, "alice.johnson@example.com", "redact off → raw email passes through")
|
||||
assert.NotContains(t, raw, "[REDACTED:", "redact off → no markers")
|
||||
}
|
||||
|
||||
func TestInvoke_NilInputAllows(t *testing.T) {
|
||||
mw := newMiddleware(t)
|
||||
out, err := mw.Invoke(context.Background(), nil)
|
||||
require.NoError(t, err, "nil input must not panic or error")
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "nil input still allows")
|
||||
assert.Empty(t, out.Metadata, "nil input emits no metadata")
|
||||
}
|
||||
Reference in New Issue
Block a user