From 6ade3839aa34de7789c66c3dfaf092ac99c99b8e Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Sat, 27 Jun 2026 00:43:35 +0200 Subject: [PATCH] [proxy] LLM parsers, pricing, and builtin middlewares (OpenAI, Anthropic, Vertex AI, AWS Bedrock) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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/internal/llm/anthropic.go | 196 ++++ proxy/internal/llm/anthropic_test.go | 169 ++++ proxy/internal/llm/bedrock.go | 189 ++++ proxy/internal/llm/bedrock_test.go | 65 ++ proxy/internal/llm/errors.go | 31 + .../llm/fixtures/anthropic_messages.json | 17 + .../llm/fixtures/anthropic_stream.txt | 21 + .../llm/fixtures/openai_chat_completion.json | 21 + .../llm/fixtures/openai_responses.json | 24 + .../llm/fixtures/openai_responses_stream.txt | 24 + proxy/internal/llm/fixtures/openai_stream.txt | 8 + proxy/internal/llm/fixtures/pricing.yaml | 59 ++ proxy/internal/llm/openai.go | 412 +++++++++ proxy/internal/llm/openai_test.go | 255 ++++++ proxy/internal/llm/parser.go | 112 +++ proxy/internal/llm/parser_test.go | 54 ++ .../llm/pricing/defaults_coverage_test.go | 65 ++ .../llm/pricing/defaults_pricing.yaml | 264 ++++++ proxy/internal/llm/pricing/pricing.go | 449 ++++++++++ proxy/internal/llm/pricing/pricing_other.go | 20 + proxy/internal/llm/pricing/pricing_test.go | 432 +++++++++ proxy/internal/llm/pricing/pricing_unix.go | 68 ++ proxy/internal/llm/sse.go | 117 +++ proxy/internal/llm/sse_test.go | 175 ++++ .../agentnetwork_chain_integration_test.go | 318 +++++++ proxy/internal/middleware/builtin/all_test.go | 40 + proxy/internal/middleware/builtin/builtin.go | 93 ++ .../middleware/builtin/cost_meter/factory.go | 88 ++ .../builtin/cost_meter/middleware.go | 193 ++++ .../builtin/cost_meter/middleware_test.go | 459 ++++++++++ .../builtin/llm_guardrail/factory.go | 82 ++ .../builtin/llm_guardrail/middleware.go | 183 ++++ .../builtin/llm_guardrail/middleware_test.go | 219 +++++ .../builtin/llm_guardrail/redact.go | 75 ++ .../builtin/llm_guardrail/redact_test.go | 217 +++++ .../builtin/llm_identity_inject/factory.go | 108 +++ .../builtin/llm_identity_inject/middleware.go | 440 +++++++++ .../llm_identity_inject/middleware_test.go | 666 ++++++++++++++ .../builtin/llm_limit_check/factory.go | 38 + .../builtin/llm_limit_check/middleware.go | 196 ++++ .../llm_limit_check/middleware_test.go | 186 ++++ .../builtin/llm_limit_record/factory.go | 35 + .../builtin/llm_limit_record/middleware.go | 144 +++ .../llm_limit_record/middleware_test.go | 191 ++++ .../llm_request_parser/bedrock_test.go | 55 ++ .../builtin/llm_request_parser/factory.go | 71 ++ .../builtin/llm_request_parser/middleware.go | 453 ++++++++++ .../llm_request_parser/middleware_test.go | 418 +++++++++ .../builtin/llm_response_parser/factory.go | 43 + .../builtin/llm_response_parser/gzip_test.go | 133 +++ .../builtin/llm_response_parser/middleware.go | 339 +++++++ .../llm_response_parser/middleware_test.go | 433 +++++++++ .../responses_stream_test.go | 69 ++ .../builtin/llm_response_parser/streaming.go | 269 ++++++ .../llm_response_parser/streaming_bedrock.go | 110 +++ .../streaming_bedrock_test.go | 74 ++ .../llm_response_parser/streaming_test.go | 169 ++++ .../middleware/builtin/llm_router/factory.go | 106 +++ .../builtin/llm_router/middleware.go | 793 +++++++++++++++++ .../builtin/llm_router/middleware_test.go | 840 ++++++++++++++++++ .../builtin/llm_router/path_routed_test.go | 159 ++++ 61 files changed, 11752 insertions(+) create mode 100644 proxy/internal/llm/anthropic.go create mode 100644 proxy/internal/llm/anthropic_test.go create mode 100644 proxy/internal/llm/bedrock.go create mode 100644 proxy/internal/llm/bedrock_test.go create mode 100644 proxy/internal/llm/errors.go create mode 100644 proxy/internal/llm/fixtures/anthropic_messages.json create mode 100644 proxy/internal/llm/fixtures/anthropic_stream.txt create mode 100644 proxy/internal/llm/fixtures/openai_chat_completion.json create mode 100644 proxy/internal/llm/fixtures/openai_responses.json create mode 100644 proxy/internal/llm/fixtures/openai_responses_stream.txt create mode 100644 proxy/internal/llm/fixtures/openai_stream.txt create mode 100644 proxy/internal/llm/fixtures/pricing.yaml create mode 100644 proxy/internal/llm/openai.go create mode 100644 proxy/internal/llm/openai_test.go create mode 100644 proxy/internal/llm/parser.go create mode 100644 proxy/internal/llm/parser_test.go create mode 100644 proxy/internal/llm/pricing/defaults_coverage_test.go create mode 100644 proxy/internal/llm/pricing/defaults_pricing.yaml create mode 100644 proxy/internal/llm/pricing/pricing.go create mode 100644 proxy/internal/llm/pricing/pricing_other.go create mode 100644 proxy/internal/llm/pricing/pricing_test.go create mode 100644 proxy/internal/llm/pricing/pricing_unix.go create mode 100644 proxy/internal/llm/sse.go create mode 100644 proxy/internal/llm/sse_test.go create mode 100644 proxy/internal/middleware/builtin/agentnetwork_chain_integration_test.go create mode 100644 proxy/internal/middleware/builtin/all_test.go create mode 100644 proxy/internal/middleware/builtin/builtin.go create mode 100644 proxy/internal/middleware/builtin/cost_meter/factory.go create mode 100644 proxy/internal/middleware/builtin/cost_meter/middleware.go create mode 100644 proxy/internal/middleware/builtin/cost_meter/middleware_test.go create mode 100644 proxy/internal/middleware/builtin/llm_guardrail/factory.go create mode 100644 proxy/internal/middleware/builtin/llm_guardrail/middleware.go create mode 100644 proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go create mode 100644 proxy/internal/middleware/builtin/llm_guardrail/redact.go create mode 100644 proxy/internal/middleware/builtin/llm_guardrail/redact_test.go create mode 100644 proxy/internal/middleware/builtin/llm_identity_inject/factory.go create mode 100644 proxy/internal/middleware/builtin/llm_identity_inject/middleware.go create mode 100644 proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go create mode 100644 proxy/internal/middleware/builtin/llm_limit_check/factory.go create mode 100644 proxy/internal/middleware/builtin/llm_limit_check/middleware.go create mode 100644 proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go create mode 100644 proxy/internal/middleware/builtin/llm_limit_record/factory.go create mode 100644 proxy/internal/middleware/builtin/llm_limit_record/middleware.go create mode 100644 proxy/internal/middleware/builtin/llm_limit_record/middleware_test.go create mode 100644 proxy/internal/middleware/builtin/llm_request_parser/bedrock_test.go create mode 100644 proxy/internal/middleware/builtin/llm_request_parser/factory.go create mode 100644 proxy/internal/middleware/builtin/llm_request_parser/middleware.go create mode 100644 proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go create mode 100644 proxy/internal/middleware/builtin/llm_response_parser/factory.go create mode 100644 proxy/internal/middleware/builtin/llm_response_parser/gzip_test.go create mode 100644 proxy/internal/middleware/builtin/llm_response_parser/middleware.go create mode 100644 proxy/internal/middleware/builtin/llm_response_parser/middleware_test.go create mode 100644 proxy/internal/middleware/builtin/llm_response_parser/responses_stream_test.go create mode 100644 proxy/internal/middleware/builtin/llm_response_parser/streaming.go create mode 100644 proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock.go create mode 100644 proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock_test.go create mode 100644 proxy/internal/middleware/builtin/llm_response_parser/streaming_test.go create mode 100644 proxy/internal/middleware/builtin/llm_router/factory.go create mode 100644 proxy/internal/middleware/builtin/llm_router/middleware.go create mode 100644 proxy/internal/middleware/builtin/llm_router/middleware_test.go create mode 100644 proxy/internal/middleware/builtin/llm_router/path_routed_test.go diff --git a/proxy/internal/llm/anthropic.go b/proxy/internal/llm/anthropic.go new file mode 100644 index 000000000..523731fbd --- /dev/null +++ b/proxy/internal/llm/anthropic.go @@ -0,0 +1,196 @@ +package llm + +import ( + "encoding/json" + "fmt" + "strings" +) + +// AnthropicParser implements the Parser interface for the Anthropic Messages +// and Completions APIs. Detection is substring-based to tolerate upstream +// path rewrites. +type AnthropicParser struct{} + +var anthropicPathHints = []string{ + "/v1/messages", + "/v1/complete", +} + +// Provider returns ProviderAnthropic. +func (AnthropicParser) Provider() Provider { return ProviderAnthropic } + +// ProviderName returns the stable label used for metrics and metadata. +func (AnthropicParser) ProviderName() string { return "anthropic" } + +// DetectFromURL reports whether the given request path looks like an +// Anthropic API endpoint. The match is case-insensitive and substring-based. +func (AnthropicParser) DetectFromURL(path string) bool { + lower := strings.ToLower(path) + for _, hint := range anthropicPathHints { + if strings.Contains(lower, hint) { + return true + } + } + return false +} + +type anthropicRequest struct { + Model string `json:"model"` + Stream *bool `json:"stream"` + System json.RawMessage `json:"system"` + Messages []anthropicMessage `json:"messages"` + // Legacy /v1/complete endpoint. + Prompt string `json:"prompt"` +} + +type anthropicMessage struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` +} + +// ParseRequest extracts the model name and streaming flag from an Anthropic +// request body. Unknown or missing fields leave the corresponding struct +// members zero-valued. +func (AnthropicParser) ParseRequest(body []byte) (RequestFacts, error) { + var req anthropicRequest + if err := json.Unmarshal(body, &req); err != nil { + return RequestFacts{}, fmt.Errorf("decode anthropic request: %w: %v", ErrMalformedRequest, err) + } + return RequestFacts{ + Model: req.Model, + Stream: ptrDeref(req.Stream), + }, nil +} + +type anthropicResponse struct { + Usage struct { + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + // CacheReadInputTokens and CacheCreationInputTokens are + // ADDITIVE to InputTokens (not subset), each billed at its + // own rate by the cost meter. cache_read is the cheaper + // read-from-cache rate, cache_creation is the more + // expensive write-to-cache rate. + CacheReadInputTokens int64 `json:"cache_read_input_tokens"` + CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"` + } `json:"usage"` +} + +// ParseResponse decodes the non-streaming Anthropic response envelope. Status +// codes other than 200 are treated as non-LLM responses so the caller can +// skip cost accounting without aborting the request. +func (AnthropicParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) { + if status != 200 { + return Usage{}, fmt.Errorf("anthropic status %d: %w", status, ErrNotLLMResponse) + } + if isEventStream(contentType) { + return Usage{}, ErrStreamingUnsupported + } + if !isJSON(contentType) { + return Usage{}, fmt.Errorf("anthropic content-type %q: %w", contentType, ErrNotLLMResponse) + } + + var resp anthropicResponse + if err := json.Unmarshal(body, &resp); err != nil { + return Usage{}, fmt.Errorf("decode anthropic response: %w: %v", ErrMalformedResponse, err) + } + return Usage{ + InputTokens: resp.Usage.InputTokens, + OutputTokens: resp.Usage.OutputTokens, + TotalTokens: resp.Usage.InputTokens + resp.Usage.OutputTokens + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens, + CachedInputTokens: resp.Usage.CacheReadInputTokens, + CacheCreationTokens: resp.Usage.CacheCreationInputTokens, + }, nil +} + +// ExtractPrompt returns the user-visible prompt text from an Anthropic +// request body. Handles the Messages API (system + messages[]) and the +// legacy /v1/complete prompt string. Returns "" on any decode failure. +func (AnthropicParser) ExtractPrompt(body []byte) string { + var req anthropicRequest + if err := json.Unmarshal(body, &req); err != nil { + return "" + } + var b strings.Builder + if len(req.System) > 0 { + if s := decodeStringOrJoin(req.System); s != "" { + b.WriteString("system: ") + b.WriteString(s) + } + } + for _, m := range req.Messages { + if b.Len() > 0 { + b.WriteByte('\n') + } + if m.Role != "" { + b.WriteString(m.Role) + b.WriteString(": ") + } + b.WriteString(decodeStringOrJoin(m.Content)) + } + if b.Len() == 0 && req.Prompt != "" { + b.WriteString(req.Prompt) + } + return b.String() +} + +// ExtractSessionID is the body-side fallback for Anthropic. Claude Code's +// authoritative session marker is the X-Claude-Code-Session-Id request +// header (handled by the request-parser middleware); this only mines the +// optional metadata.user_id for an embedded "...session_" marker. +// metadata.user_id on its own is a USER identifier, not a session, so the +// whole value is deliberately NOT used — returning it would mislabel every +// request from a user as one session. Returns "" when no session marker is +// present. +func (AnthropicParser) ExtractSessionID(body []byte) string { + var req struct { + Metadata struct { + UserID string `json:"user_id"` + } `json:"metadata"` + } + if err := json.Unmarshal(body, &req); err != nil { + return "" + } + if idx := strings.LastIndex(req.Metadata.UserID, "session_"); idx >= 0 { + if session := req.Metadata.UserID[idx+len("session_"):]; session != "" { + return session + } + } + return "" +} + +type anthropicMessageResponse struct { + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"content"` + // Legacy /v1/complete response. + Completion string `json:"completion"` +} + +// ExtractCompletion returns the assistant text from a non-streaming Anthropic +// Messages or Completions response. Returns "" when status/content-type +// indicate the body is not parseable or no text part is present. +func (AnthropicParser) ExtractCompletion(status int, contentType string, body []byte) string { + if status != 200 || isEventStream(contentType) || !isJSON(contentType) { + return "" + } + var resp anthropicMessageResponse + if err := json.Unmarshal(body, &resp); err != nil { + return "" + } + var b strings.Builder + for _, part := range resp.Content { + if part.Text == "" { + continue + } + if b.Len() > 0 { + b.WriteByte('\n') + } + b.WriteString(part.Text) + } + if b.Len() == 0 { + return resp.Completion + } + return b.String() +} diff --git a/proxy/internal/llm/anthropic_test.go b/proxy/internal/llm/anthropic_test.go new file mode 100644 index 000000000..a0c1f7896 --- /dev/null +++ b/proxy/internal/llm/anthropic_test.go @@ -0,0 +1,169 @@ +package llm + +import ( + "errors" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAnthropicDetectFromURL(t *testing.T) { + p := AnthropicParser{} + + cases := map[string]bool{ + "/v1/messages": true, + "/v1/complete": true, + "/V1/Messages": true, + "/proxy/v1/messages?x": true, + "/v1/chat/completions": false, + "": false, + } + for path, want := range cases { + assert.Equal(t, want, p.DetectFromURL(path), "DetectFromURL(%q)", path) + } +} + +func TestAnthropicParseRequest(t *testing.T) { + p := AnthropicParser{} + + t.Run("stream true", func(t *testing.T) { + facts, err := p.ParseRequest([]byte(`{"model":"claude-sonnet-4-5","stream":true}`)) + require.NoError(t, err) + assert.Equal(t, "claude-sonnet-4-5", facts.Model, "model extracted") + assert.True(t, facts.Stream, "stream flag honoured") + }) + + t.Run("stream default", func(t *testing.T) { + facts, err := p.ParseRequest([]byte(`{"model":"claude-sonnet-4-5"}`)) + require.NoError(t, err) + assert.False(t, facts.Stream, "missing stream flag defaults to false") + }) + + t.Run("malformed", func(t *testing.T) { + _, err := p.ParseRequest([]byte(`{"model":`)) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrMalformedRequest), "sentinel wrapped") + }) +} + +func TestAnthropicParseResponse(t *testing.T) { + p := AnthropicParser{} + + t.Run("happy fixture", func(t *testing.T) { + body, err := os.ReadFile(filepath.Join("fixtures", "anthropic_messages.json")) + require.NoError(t, err, "fixture must be readable") + + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Equal(t, int64(123), usage.InputTokens, "input tokens extracted") + assert.Equal(t, int64(45), usage.OutputTokens, "output tokens extracted") + assert.Equal(t, int64(168), usage.TotalTokens, "total computed as sum") + }) + + t.Run("streaming rejected", func(t *testing.T) { + _, err := p.ParseResponse(200, "text/event-stream", []byte("")) + require.ErrorIs(t, err, ErrStreamingUnsupported, "SSE responses must use the scanner") + }) + + t.Run("non-200", func(t *testing.T) { + _, err := p.ParseResponse(429, "application/json", []byte(`{}`)) + require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 rejected as non-LLM") + }) + + t.Run("non-json content type", func(t *testing.T) { + _, err := p.ParseResponse(200, "text/html", []byte(`{}`)) + require.ErrorIs(t, err, ErrNotLLMResponse, "text/html treated as non-LLM") + }) + + t.Run("malformed body", func(t *testing.T) { + _, err := p.ParseResponse(200, "application/json", []byte(`{`)) + require.ErrorIs(t, err, ErrMalformedResponse, "bad JSON yields malformed error") + }) + + // Anthropic's two cache fields are ADDITIVE to input_tokens (not + // subset). The parser must surface them so the cost meter can + // bill each bucket at its own configured rate. Total includes + // every bucket so downstream attribution sees the full token + // volume the request consumed. + t.Run("cache_read_input_tokens surfaces as CachedInputTokens (additive)", func(t *testing.T) { + body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768}}`) + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Equal(t, int64(256), usage.InputTokens, "regular input remains separate from cache buckets") + assert.Equal(t, int64(768), usage.CachedInputTokens, "cache_read maps onto CachedInputTokens — same field carries OpenAI cached subset and Anthropic cache reads") + assert.Zero(t, usage.CacheCreationTokens) + assert.Equal(t, int64(256+200+768), usage.TotalTokens, "total includes every input bucket plus output — cache reads are billable tokens") + }) + + t.Run("cache_creation_input_tokens surfaces as CacheCreationTokens (additive)", func(t *testing.T) { + body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_creation_input_tokens":512}}`) + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Equal(t, int64(256), usage.InputTokens) + assert.Zero(t, usage.CachedInputTokens) + assert.Equal(t, int64(512), usage.CacheCreationTokens, "cache_creation surfaces — meter applies the write-rate multiplier") + assert.Equal(t, int64(256+200+512), usage.TotalTokens) + }) + + t.Run("both cache buckets present", func(t *testing.T) { + body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768,"cache_creation_input_tokens":512}}`) + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Equal(t, int64(768), usage.CachedInputTokens) + assert.Equal(t, int64(512), usage.CacheCreationTokens) + assert.Equal(t, int64(256+200+768+512), usage.TotalTokens, "all four buckets sum into total") + }) + + t.Run("absent cache fields leave counts at zero", func(t *testing.T) { + body := []byte(`{"usage":{"input_tokens":100,"output_tokens":50}}`) + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Zero(t, usage.CachedInputTokens, "no cache_read field = no cached count") + assert.Zero(t, usage.CacheCreationTokens, "no cache_creation field = no creation count") + assert.Equal(t, int64(150), usage.TotalTokens, "back to the simple in+out total when no cache buckets present") + }) +} + +func TestAnthropicExtractPrompt_Messages(t *testing.T) { + body := []byte(`{"model":"claude-sonnet-4-7","system":"be brief","messages":[{"role":"user","content":"hi"},{"role":"assistant","content":"yes"}]}`) + got := AnthropicParser{}.ExtractPrompt(body) + require.Contains(t, got, "system: be brief", "system surfaces with role label") + require.Contains(t, got, "user: hi", "user message surfaces") + require.Contains(t, got, "assistant: yes", "assistant message surfaces") +} + +func TestAnthropicExtractPrompt_LegacyComplete(t *testing.T) { + body := []byte(`{"model":"claude-2","prompt":"\n\nHuman: hi\n\nAssistant:"}`) + got := AnthropicParser{}.ExtractPrompt(body) + require.Contains(t, got, "Human: hi", "legacy prompt string surfaces") +} + +func TestAnthropicExtractSessionID(t *testing.T) { + t.Run("claude code session suffix", func(t *testing.T) { + body := []byte(`{"model":"claude-opus-4-8","metadata":{"user_id":"user_abc123_account_def456_session_9f8e7d6c"},"messages":[]}`) + assert.Equal(t, "9f8e7d6c", AnthropicParser{}.ExtractSessionID(body), "session_ suffix must be extracted from metadata.user_id") + }) + t.Run("plain user_id is not treated as a session", func(t *testing.T) { + body := []byte(`{"model":"claude-opus-4-8","metadata":{"user_id":"acme-team"},"messages":[]}`) + assert.Equal(t, "", AnthropicParser{}.ExtractSessionID(body), "a user identifier without a session marker must NOT be used as a session id") + }) + t.Run("no metadata yields empty", func(t *testing.T) { + body := []byte(`{"model":"claude-opus-4-8","messages":[{"role":"user","content":"hi"}]}`) + assert.Equal(t, "", AnthropicParser{}.ExtractSessionID(body), "absent metadata.user_id yields no session id") + }) +} + +func TestAnthropicExtractCompletion_Messages(t *testing.T) { + body, err := os.ReadFile(filepath.Join("fixtures", "anthropic_messages.json")) + require.NoError(t, err) + got := AnthropicParser{}.ExtractCompletion(200, "application/json", body) + require.NotEmpty(t, got, "anthropic fixture has assistant text") +} + +func TestAnthropicExtractCompletion_Streaming(t *testing.T) { + got := AnthropicParser{}.ExtractCompletion(200, "text/event-stream", []byte("")) + require.Empty(t, got, "streaming responses are skipped") +} diff --git a/proxy/internal/llm/bedrock.go b/proxy/internal/llm/bedrock.go new file mode 100644 index 000000000..f7802beb2 --- /dev/null +++ b/proxy/internal/llm/bedrock.go @@ -0,0 +1,189 @@ +package llm + +import ( + "encoding/json" + "fmt" + "strings" +) + +// ProviderNameBedrock is the stable label for the AWS Bedrock parser, used as +// the llm.provider metadata value and the cost-meter formula selector. +const ProviderNameBedrock = "bedrock" + +// BedrockParser implements the Parser interface for the AWS Bedrock runtime. +// Bedrock carries the model in the URL path (/model/{id}/{action}); the request +// middleware extracts it there, so this parser focuses on the response shapes: +// the vendor-native InvokeModel body (e.g. Anthropic's snake_case usage) and the +// unified Converse body (camelCase usage). +type BedrockParser struct{} + +var bedrockPathHints = []string{"/invoke", "/converse"} + +// Provider returns ProviderBedrock. +func (BedrockParser) Provider() Provider { return ProviderBedrock } + +// ProviderName returns the stable label used for metrics and metadata. +func (BedrockParser) ProviderName() string { return ProviderNameBedrock } + +// DetectFromURL reports whether the path is a Bedrock runtime model endpoint. +func (BedrockParser) DetectFromURL(path string) bool { + lower := strings.ToLower(path) + if !strings.HasPrefix(lower, "/model/") { + return false + } + for _, hint := range bedrockPathHints { + if strings.Contains(lower, hint) { + return true + } + } + return false +} + +// ParseRequest is a no-op for Bedrock: the model lives in the URL path, not the +// body, and the streaming flag is derived from the path action. The request +// middleware handles both via parseBedrockPath, so this returns empty facts. +func (BedrockParser) ParseRequest([]byte) (RequestFacts, error) { + return RequestFacts{}, nil +} + +// bedrockResponse captures token usage from both Bedrock response shapes: +// InvokeModel (vendor-native; Anthropic uses snake_case + additive cache +// buckets) and Converse (camelCase, with a precomputed total). +type bedrockResponse struct { + Usage struct { + // InvokeModel (Anthropic-on-Bedrock) — snake_case. + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + CacheReadInputTokens int64 `json:"cache_read_input_tokens"` + CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"` + // Converse — camelCase. + InputTokensCamel int64 `json:"inputTokens"` + OutputTokensCamel int64 `json:"outputTokens"` + TotalTokensCamel int64 `json:"totalTokens"` + } `json:"usage"` +} + +// ParseResponse decodes the non-streaming Bedrock response envelope, handling +// both the InvokeModel and Converse usage shapes. Non-200 / non-JSON bodies are +// treated as non-LLM responses so the caller skips cost accounting. +func (BedrockParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) { + if status != 200 { + return Usage{}, fmt.Errorf("bedrock status %d: %w", status, ErrNotLLMResponse) + } + if isAWSEventStream(contentType) || isEventStream(contentType) { + return Usage{}, ErrStreamingUnsupported + } + if !isJSON(contentType) { + return Usage{}, fmt.Errorf("bedrock content-type %q: %w", contentType, ErrNotLLMResponse) + } + + var resp bedrockResponse + if err := json.Unmarshal(body, &resp); err != nil { + return Usage{}, fmt.Errorf("decode bedrock response: %w: %v", ErrMalformedResponse, err) + } + inTok := firstNonZero(resp.Usage.InputTokens, resp.Usage.InputTokensCamel) + outTok := firstNonZero(resp.Usage.OutputTokens, resp.Usage.OutputTokensCamel) + total := resp.Usage.TotalTokensCamel + if total == 0 { + total = inTok + outTok + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens + } + return Usage{ + InputTokens: inTok, + OutputTokens: outTok, + TotalTokens: total, + CachedInputTokens: resp.Usage.CacheReadInputTokens, + CacheCreationTokens: resp.Usage.CacheCreationInputTokens, + }, nil +} + +// ExtractPrompt returns the user-visible prompt from a Bedrock request body, +// handling both the InvokeModel (Anthropic Messages: system + messages[]) and +// Converse (messages[].content[].text) shapes. Returns "" on decode failure. +func (BedrockParser) ExtractPrompt(body []byte) string { + var req struct { + System json.RawMessage `json:"system"` + Messages []struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` + } `json:"messages"` + } + if err := json.Unmarshal(body, &req); err != nil { + return "" + } + var b strings.Builder + if s := decodeStringOrJoin(req.System); s != "" { + b.WriteString("system: ") + b.WriteString(s) + } + for _, m := range req.Messages { + if b.Len() > 0 { + b.WriteByte('\n') + } + if m.Role != "" { + b.WriteString(m.Role) + b.WriteString(": ") + } + b.WriteString(decodeStringOrJoin(m.Content)) + } + return b.String() +} + +// ExtractCompletion returns the assistant text from a non-streaming Bedrock +// response, handling InvokeModel (Anthropic content[].text) and Converse +// (output.message.content[].text). +func (BedrockParser) ExtractCompletion(status int, contentType string, body []byte) string { + if status != 200 || isAWSEventStream(contentType) || !isJSON(contentType) { + return "" + } + var resp struct { + Content []struct { + Text string `json:"text"` + } `json:"content"` + Output struct { + Message struct { + Content []struct { + Text string `json:"text"` + } `json:"content"` + } `json:"message"` + } `json:"output"` + } + if err := json.Unmarshal(body, &resp); err != nil { + return "" + } + var b strings.Builder + appendText := func(text string) { + if text == "" { + return + } + if b.Len() > 0 { + b.WriteByte('\n') + } + b.WriteString(text) + } + for _, p := range resp.Content { + appendText(p.Text) + } + for _, p := range resp.Output.Message.Content { + appendText(p.Text) + } + return b.String() +} + +// ExtractSessionID has no Bedrock-native marker; session grouping relies on the +// request headers handled by the middleware. Returns "". +func (BedrockParser) ExtractSessionID([]byte) string { return "" } + +// firstNonZero returns a when non-zero, else b. Folds the snake_case and +// camelCase usage variants into a single value. +func firstNonZero(a, b int64) int64 { + if a != 0 { + return a + } + return b +} + +// isAWSEventStream reports whether contentType is the AWS binary event-stream +// framing used by Bedrock's streaming endpoints. +func isAWSEventStream(contentType string) bool { + return strings.Contains(strings.ToLower(contentType), "application/vnd.amazon.eventstream") +} diff --git a/proxy/internal/llm/bedrock_test.go b/proxy/internal/llm/bedrock_test.go new file mode 100644 index 000000000..ca6f092f3 --- /dev/null +++ b/proxy/internal/llm/bedrock_test.go @@ -0,0 +1,65 @@ +package llm + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestBedrockParser_ParseResponse_Invoke(t *testing.T) { + body := []byte(`{"usage":{"input_tokens":13,"output_tokens":5,"cache_read_input_tokens":2,"cache_creation_input_tokens":4}}`) + u, err := BedrockParser{}.ParseResponse(200, "application/json", body) + require.NoError(t, err) + require.Equal(t, int64(13), u.InputTokens, "invoke input tokens") + require.Equal(t, int64(5), u.OutputTokens, "invoke output tokens") + require.Equal(t, int64(2), u.CachedInputTokens, "invoke cache-read tokens") + require.Equal(t, int64(4), u.CacheCreationTokens, "invoke cache-creation tokens") + require.Equal(t, int64(13+5+2+4), u.TotalTokens, "invoke total is additive") +} + +func TestBedrockParser_ParseResponse_Converse(t *testing.T) { + body := []byte(`{"output":{"message":{"content":[{"text":"pong"}]}},"usage":{"inputTokens":11,"outputTokens":3,"totalTokens":14}}`) + u, err := BedrockParser{}.ParseResponse(200, "application/json", body) + require.NoError(t, err) + require.Equal(t, int64(11), u.InputTokens, "converse camelCase input tokens") + require.Equal(t, int64(3), u.OutputTokens, "converse camelCase output tokens") + require.Equal(t, int64(14), u.TotalTokens, "converse uses provider total") +} + +func TestBedrockParser_ParseResponse_StreamingUnsupported(t *testing.T) { + _, err := BedrockParser{}.ParseResponse(200, "application/vnd.amazon.eventstream", []byte("binary")) + require.ErrorIs(t, err, ErrStreamingUnsupported, "event-stream must route to the streaming accumulator") +} + +func TestBedrockParser_ParseResponse_NonSuccess(t *testing.T) { + _, err := BedrockParser{}.ParseResponse(404, "application/json", []byte(`{"message":"gated"}`)) + require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 is not an LLM response") +} + +func TestBedrockParser_ExtractCompletion(t *testing.T) { + invoke := BedrockParser{}.ExtractCompletion(200, "application/json", []byte(`{"content":[{"text":"a"},{"text":"b"}]}`)) + require.Equal(t, "a\nb", invoke, "invoke completion joins content parts") + + converse := BedrockParser{}.ExtractCompletion(200, "application/json", []byte(`{"output":{"message":{"content":[{"text":"x"}]}}}`)) + require.Equal(t, "x", converse, "converse completion reads output.message.content") +} + +func TestBedrockParser_ExtractPrompt(t *testing.T) { + invoke := BedrockParser{}.ExtractPrompt([]byte(`{"messages":[{"role":"user","content":"hi"}]}`)) + require.Equal(t, "user: hi", invoke, "invoke prompt reads anthropic content string") + + converse := BedrockParser{}.ExtractPrompt([]byte(`{"messages":[{"role":"user","content":[{"text":"hello"}]}]}`)) + require.Equal(t, "user: hello", converse, "converse prompt reads content parts") +} + +func TestBedrockParser_DetectFromURL(t *testing.T) { + require.True(t, BedrockParser{}.DetectFromURL("/model/eu.anthropic.claude/invoke"), "invoke path") + require.True(t, BedrockParser{}.DetectFromURL("/model/x/converse-stream"), "converse-stream path") + require.False(t, BedrockParser{}.DetectFromURL("/v1/chat/completions"), "openai path is not bedrock") +} + +func TestBedrockParser_RegisteredByName(t *testing.T) { + p, ok := ParserByName(ProviderNameBedrock) + require.True(t, ok, "bedrock parser is registered") + require.Equal(t, ProviderNameBedrock, p.ProviderName()) +} diff --git a/proxy/internal/llm/errors.go b/proxy/internal/llm/errors.go new file mode 100644 index 000000000..09019fbbe --- /dev/null +++ b/proxy/internal/llm/errors.go @@ -0,0 +1,31 @@ +package llm + +import "errors" + +// Sentinel errors returned by parsers and the pricing loader. Callers use +// errors.Is to branch on a condition without coupling to parser internals. +var ( + // ErrUnknownProvider indicates no parser claimed the request path. + ErrUnknownProvider = errors.New("llmobs: unknown provider") + + // ErrUnsupportedModel indicates the response parsed successfully but the + // model is absent from the pricing table. Token counts are still valid. + ErrUnsupportedModel = errors.New("llmobs: unsupported model") + + // ErrNotLLMResponse indicates the response is not a JSON success body + // that a non-streaming parser can consume (non-200 or wrong content type). + ErrNotLLMResponse = errors.New("llmobs: not an LLM response") + + // ErrStreamingUnsupported indicates the caller passed an SSE response to + // a non-streaming parser. Streaming is handled separately via the SSE + // scanner. + ErrStreamingUnsupported = errors.New("llmobs: streaming response requires SSE scanner") + + // ErrMalformedResponse indicates the response body could not be decoded + // as the provider-specific JSON schema. + ErrMalformedResponse = errors.New("llmobs: malformed response body") + + // ErrMalformedRequest indicates the request body could not be decoded as + // the provider-specific JSON schema. + ErrMalformedRequest = errors.New("llmobs: malformed request body") +) diff --git a/proxy/internal/llm/fixtures/anthropic_messages.json b/proxy/internal/llm/fixtures/anthropic_messages.json new file mode 100644 index 000000000..2c9bb663a --- /dev/null +++ b/proxy/internal/llm/fixtures/anthropic_messages.json @@ -0,0 +1,17 @@ +{ + "id": "msg_abc", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [ + { + "type": "text", + "text": "Hello, world!" + } + ], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 123, + "output_tokens": 45 + } +} diff --git a/proxy/internal/llm/fixtures/anthropic_stream.txt b/proxy/internal/llm/fixtures/anthropic_stream.txt new file mode 100644 index 000000000..2b8bb889c --- /dev/null +++ b/proxy/internal/llm/fixtures/anthropic_stream.txt @@ -0,0 +1,21 @@ +event: message_start +data: {"type":"message_start","message":{"id":"msg_abc","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[],"stop_reason":null,"usage":{"input_tokens":123,"output_tokens":1}}} + +event: content_block_start +data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":", world!"}} + +event: content_block_stop +data: {"type":"content_block_stop","index":0} + +event: message_delta +data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":45}} + +event: message_stop +data: {"type":"message_stop"} + diff --git a/proxy/internal/llm/fixtures/openai_chat_completion.json b/proxy/internal/llm/fixtures/openai_chat_completion.json new file mode 100644 index 000000000..d0e25337b --- /dev/null +++ b/proxy/internal/llm/fixtures/openai_chat_completion.json @@ -0,0 +1,21 @@ +{ + "id": "chatcmpl-abc", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello, world!" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 123, + "completion_tokens": 45, + "total_tokens": 168 + } +} diff --git a/proxy/internal/llm/fixtures/openai_responses.json b/proxy/internal/llm/fixtures/openai_responses.json new file mode 100644 index 000000000..f998fcd33 --- /dev/null +++ b/proxy/internal/llm/fixtures/openai_responses.json @@ -0,0 +1,24 @@ +{ + "id": "resp_abc", + "object": "response", + "created_at": 1700000000, + "model": "gpt-5.4", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok"}] + } + ], + "usage": { + "input_tokens": 15, + "input_tokens_details": { + "cached_tokens": 0 + }, + "output_tokens": 414, + "output_tokens_details": { + "reasoning_tokens": 0 + }, + "total_tokens": 429 + } +} diff --git a/proxy/internal/llm/fixtures/openai_responses_stream.txt b/proxy/internal/llm/fixtures/openai_responses_stream.txt new file mode 100644 index 000000000..2801fa99d --- /dev/null +++ b/proxy/internal/llm/fixtures/openai_responses_stream.txt @@ -0,0 +1,24 @@ +event: response.created +data: {"type":"response.created","response":{"id":"resp_abc","object":"response","model":"gpt-5.5","usage":null}} + +event: response.in_progress +data: {"type":"response.in_progress","response":{"id":"resp_abc","usage":null}} + +event: response.output_item.added +data: {"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant","content":[]}} + +event: response.content_part.added +data: {"type":"response.content_part.added","item_id":"msg_1","output_index":0,"content_index":0,"part":{"type":"output_text","text":""}} + +event: response.output_text.delta +data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"Hello"} + +event: response.output_text.delta +data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":", world!"} + +event: response.output_text.done +data: {"type":"response.output_text.done","item_id":"msg_1","output_index":0,"content_index":0,"text":"Hello, world!"} + +event: response.completed +data: {"type":"response.completed","response":{"id":"resp_abc","object":"response","model":"gpt-5.5","usage":{"input_tokens":123,"input_tokens_details":{"cached_tokens":40},"output_tokens":45,"output_tokens_details":{"reasoning_tokens":12},"total_tokens":168}}} + diff --git a/proxy/internal/llm/fixtures/openai_stream.txt b/proxy/internal/llm/fixtures/openai_stream.txt new file mode 100644 index 000000000..058b7ce22 --- /dev/null +++ b/proxy/internal/llm/fixtures/openai_stream.txt @@ -0,0 +1,8 @@ +data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]} + +data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":", world!"},"finish_reason":null}]} + +data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":123,"completion_tokens":45,"total_tokens":168}} + +data: [DONE] + diff --git a/proxy/internal/llm/fixtures/pricing.yaml b/proxy/internal/llm/fixtures/pricing.yaml new file mode 100644 index 000000000..3d26ff803 --- /dev/null +++ b/proxy/internal/llm/fixtures/pricing.yaml @@ -0,0 +1,59 @@ +# Realistic-pricing starter for llm_observability. Drop this into the +# directory you point the proxy at via --plugin-data-dir, then reference it +# from the target's plugin config: +# +# plugins: +# - id: llm_observability +# enabled: true +# params: +# pricing_path: pricing.yaml +# +# Values are USD per 1_000 tokens. Public list prices drift; treat this as a +# starting point and keep your production copy current. + +openai: + # GPT-5 family + gpt-5: + input_per_1k: 0.00125 + output_per_1k: 0.01 + gpt-5-mini: + input_per_1k: 0.00025 + output_per_1k: 0.002 + gpt-5-nano: + input_per_1k: 0.00005 + output_per_1k: 0.0004 + gpt-5.4: + input_per_1k: 0.00125 + output_per_1k: 0.01 + # GPT-4o family + gpt-4o: + input_per_1k: 0.0025 + output_per_1k: 0.01 + gpt-4o-mini: + input_per_1k: 0.00015 + output_per_1k: 0.0006 + # Embeddings + text-embedding-3-large: + input_per_1k: 0.00013 + output_per_1k: 0 + text-embedding-3-small: + input_per_1k: 0.00002 + output_per_1k: 0 + +anthropic: + # Claude 4.x family + claude-opus-4-7: + input_per_1k: 0.015 + output_per_1k: 0.075 + claude-sonnet-4-7: + input_per_1k: 0.003 + output_per_1k: 0.015 + claude-sonnet-4-6: + input_per_1k: 0.003 + output_per_1k: 0.015 + claude-sonnet-4-5: + input_per_1k: 0.003 + output_per_1k: 0.015 + claude-haiku-4-5: + input_per_1k: 0.0008 + output_per_1k: 0.004 diff --git a/proxy/internal/llm/openai.go b/proxy/internal/llm/openai.go new file mode 100644 index 000000000..86ee30797 --- /dev/null +++ b/proxy/internal/llm/openai.go @@ -0,0 +1,412 @@ +package llm + +import ( + "encoding/json" + "fmt" + "strings" +) + +// OpenAIParser implements the Parser interface for OpenAI-compatible APIs. +// It recognizes chat.completions, completions, embeddings, and the newer +// responses endpoint; any proxy path-prefix stripping is tolerated by the +// substring match in DetectFromURL. +type OpenAIParser struct{} + +// openAIPathHints are substring patterns that mark a request as +// OpenAI-shaped. The bare `/chat/completions` is listed alongside +// `/v1/chat/completions` because gateways like Cloudflare AI +// Gateway place their own version segment before the provider +// slug (gateway/v1/{account}/{gateway}/openai/chat/completions) — +// the canonical `/v1/` ends up nowhere near `/chat/completions`, +// so the `/v1/chat/completions` hint misses. `/chat/completions` +// is OpenAI's API contract: any service accepting OpenAI bodies +// serves at this path, so false-positive risk is negligible. +// `/completions` (legacy), `/embeddings`, and `/responses` are +// kept on the canonical-only path because their bare forms are +// too generic to be safe substrings. +var openAIPathHints = []string{ + "/v1/chat/completions", + "/v1/completions", + "/v1/embeddings", + "/v1/responses", + "/chat/completions", +} + +// Provider returns ProviderOpenAI. +func (OpenAIParser) Provider() Provider { return ProviderOpenAI } + +// ProviderName returns the stable label used for metrics and metadata. +func (OpenAIParser) ProviderName() string { return "openai" } + +// DetectFromURL reports whether the given request path looks like an OpenAI +// API endpoint. The match is case-insensitive and substring-based so that a +// reverse proxy prefix strip or rewrite does not defeat detection. +func (OpenAIParser) DetectFromURL(path string) bool { + lower := strings.ToLower(path) + for _, hint := range openAIPathHints { + if strings.Contains(lower, hint) { + return true + } + } + return false +} + +type openAIRequest struct { + Model string `json:"model"` + Stream *bool `json:"stream"` + StreamOptions *struct { + IncludeUsage *bool `json:"include_usage"` + } `json:"stream_options"` + // Chat Completions / Completions: messages[].content (string or array of + // content parts). Responses API: input is either a string or an array of + // items with content parts. We use json.RawMessage to defer parsing each + // shape independently. + Messages []openAIMessage `json:"messages"` + Prompt json.RawMessage `json:"prompt"` + Input json.RawMessage `json:"input"` +} + +type openAIMessage struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` +} + +// ParseRequest extracts the model name and streaming flag from an OpenAI +// request body. Unknown or missing fields leave the corresponding struct +// members zero-valued. +func (OpenAIParser) ParseRequest(body []byte) (RequestFacts, error) { + var req openAIRequest + if err := json.Unmarshal(body, &req); err != nil { + return RequestFacts{}, fmt.Errorf("decode openai request: %w: %v", ErrMalformedRequest, err) + } + return RequestFacts{ + Model: req.Model, + Stream: ptrDeref(req.Stream), + }, nil +} + +// openAIResponse accepts both naming conventions in a single struct because +// OpenAI's older Chat Completions API uses prompt_tokens/completion_tokens +// while the newer Responses API (/v1/responses) uses input_tokens/output_tokens +// (aligned with Anthropic). Pointer fields let us tell "absent" from "zero". +// +// PromptTokensDetails.CachedTokens (Chat Completions) and +// InputTokensDetails.CachedTokens (Responses API) carry the SUBSET of +// prompt/input tokens that hit the prompt cache. Cost-meter applies the +// discount rate to that subset and the regular rate to the remainder so +// we never double-bill the cached portion. +type openAIResponse struct { + Usage struct { + PromptTokens *int64 `json:"prompt_tokens"` + CompletionTokens *int64 `json:"completion_tokens"` + InputTokens *int64 `json:"input_tokens"` + OutputTokens *int64 `json:"output_tokens"` + TotalTokens *int64 `json:"total_tokens"` + PromptTokensDetails *struct { + CachedTokens *int64 `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + InputTokensDetails *struct { + CachedTokens *int64 `json:"cached_tokens"` + } `json:"input_tokens_details"` + } `json:"usage"` +} + +// ParseResponse decodes the non-streaming OpenAI response envelope. Status +// codes other than 200 are treated as non-LLM responses so the caller can +// skip cost accounting without aborting the request. +func (OpenAIParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) { + if status != 200 { + return Usage{}, fmt.Errorf("openai status %d: %w", status, ErrNotLLMResponse) + } + if isEventStream(contentType) { + return Usage{}, ErrStreamingUnsupported + } + if !isJSON(contentType) { + return Usage{}, fmt.Errorf("openai content-type %q: %w", contentType, ErrNotLLMResponse) + } + + var resp openAIResponse + if err := json.Unmarshal(body, &resp); err != nil { + return Usage{}, fmt.Errorf("decode openai response: %w: %v", ErrMalformedResponse, err) + } + + // Responses-API names take precedence when present; fall back to the older + // Chat Completions names. This handles both endpoints transparently + // without forcing a per-route configuration. + u := Usage{ + InputTokens: pickInt64(resp.Usage.InputTokens, resp.Usage.PromptTokens), + OutputTokens: pickInt64(resp.Usage.OutputTokens, resp.Usage.CompletionTokens), + TotalTokens: derefInt64(resp.Usage.TotalTokens), + CachedInputTokens: openAICachedTokens(resp), + } + if u.TotalTokens == 0 && (u.InputTokens > 0 || u.OutputTokens > 0) { + u.TotalTokens = u.InputTokens + u.OutputTokens + } + return u, nil +} + +// openAICachedTokens returns the cached-prompt subset reported by +// either the Responses-API (input_tokens_details.cached_tokens) or +// the Chat-Completions API (prompt_tokens_details.cached_tokens). +// Responses-API takes precedence when both are populated. +func openAICachedTokens(resp openAIResponse) int64 { + // Responses-API details are authoritative when present: an explicit + // cached_tokens of 0 must be honored, not treated as missing and + // overridden by the Chat-Completions field (which would overstate cache). + if resp.Usage.InputTokensDetails != nil && resp.Usage.InputTokensDetails.CachedTokens != nil { + return derefInt64(resp.Usage.InputTokensDetails.CachedTokens) + } + if resp.Usage.PromptTokensDetails != nil { + return derefInt64(resp.Usage.PromptTokensDetails.CachedTokens) + } + return 0 +} + +// ExtractPrompt returns the user-visible prompt text from an OpenAI request. +// Handles chat.completions (messages[].content), legacy completions (prompt +// string), and the Responses API (input as string or content-part array). +// Returns "" when nothing extractable is found. +func (OpenAIParser) ExtractPrompt(body []byte) string { + var req openAIRequest + if err := json.Unmarshal(body, &req); err != nil { + return "" + } + if len(req.Messages) > 0 { + return joinMessages(req.Messages) + } + if len(req.Input) > 0 { + return extractResponsesInput(req.Input) + } + if len(req.Prompt) > 0 { + return decodeStringOrJoin(req.Prompt) + } + return "" +} + +// extractResponsesInput handles the Responses API `input` field. It is one +// of three shapes: a plain string, an array of message items +// ({role, content: string | [parts]}) as sent by Codex and the Responses +// SDK, or a flat array of content parts ({type, text/input_text}). Message +// items are flattened to "role: text" lines; items without extractable text +// (reasoning blocks, tool calls) are skipped. +func extractResponsesInput(raw json.RawMessage) string { + if s, ok := tryDecodeString(raw); ok { + return s + } + var items []struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` + Text string `json:"text"` + InputText string `json:"input_text"` + } + if err := json.Unmarshal(raw, &items); err != nil { + return extractContentParts(raw) + } + var b strings.Builder + for _, it := range items { + var text string + switch { + case len(it.Content) > 0: + text = decodeStringOrJoin(it.Content) + case it.Text != "": + text = it.Text + case it.InputText != "": + text = it.InputText + } + if text == "" { + continue + } + if b.Len() > 0 { + b.WriteByte('\n') + } + if it.Role != "" { + b.WriteString(it.Role) + b.WriteString(": ") + } + b.WriteString(text) + } + return b.String() +} + +// ExtractSessionID reads the OpenAI session marker. Codex (the Responses +// API client) stamps client_metadata.session_id on every request body; +// plain chat.completions traffic carries no session id and yields "". +func (OpenAIParser) ExtractSessionID(body []byte) string { + var req struct { + ClientMetadata struct { + SessionID string `json:"session_id"` + } `json:"client_metadata"` + } + if err := json.Unmarshal(body, &req); err != nil { + return "" + } + return req.ClientMetadata.SessionID +} + +type openAIChatChoice struct { + Message struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` + } `json:"message"` + Text string `json:"text"` +} + +type openAIChatResponse struct { + Choices []openAIChatChoice `json:"choices"` + // Responses API: output[].content[].text + Output []struct { + Type string `json:"type"` + Content json.RawMessage `json:"content"` + Text string `json:"text"` + } `json:"output"` + OutputText string `json:"output_text"` +} + +// ExtractCompletion returns the assistant text from a non-streaming OpenAI +// response. Handles chat.completions (choices[].message.content), legacy +// completions (choices[].text), and Responses API (output[].content[].text +// or the convenience output_text field). +func (OpenAIParser) ExtractCompletion(status int, contentType string, body []byte) string { + if status != 200 || isEventStream(contentType) || !isJSON(contentType) { + return "" + } + var resp openAIChatResponse + if err := json.Unmarshal(body, &resp); err != nil { + return "" + } + if resp.OutputText != "" { + return resp.OutputText + } + for _, c := range resp.Choices { + if len(c.Message.Content) > 0 { + if s := decodeStringOrJoin(c.Message.Content); s != "" { + return s + } + } + if c.Text != "" { + return c.Text + } + } + for _, o := range resp.Output { + if o.Text != "" { + return o.Text + } + if len(o.Content) > 0 { + if s := extractContentParts(o.Content); s != "" { + return s + } + } + } + return "" +} + +// joinMessages flattens a chat.completions messages array into a single +// "role: content" string per message, separated by newlines. Roles surface +// system/user/assistant context which is useful for log review. +func joinMessages(msgs []openAIMessage) string { + var b strings.Builder + for i, m := range msgs { + if i > 0 { + b.WriteByte('\n') + } + if m.Role != "" { + b.WriteString(m.Role) + b.WriteString(": ") + } + b.WriteString(decodeStringOrJoin(m.Content)) + } + return b.String() +} + +// extractContentParts handles the Responses-API content shape, which is +// either a single string or an array of {type, text} parts. text and +// input_text both carry user-facing content. +func extractContentParts(raw json.RawMessage) string { + if s, ok := tryDecodeString(raw); ok { + return s + } + var parts []struct { + Type string `json:"type"` + Text string `json:"text"` + InputText string `json:"input_text"` + } + if err := json.Unmarshal(raw, &parts); err != nil { + // Last-ditch: array of strings. + var arr []string + if json.Unmarshal(raw, &arr) == nil { + return strings.Join(arr, "\n") + } + return "" + } + var b strings.Builder + for _, p := range parts { + var text string + switch { + case p.Text != "": + text = p.Text + case p.InputText != "": + text = p.InputText + } + if text == "" { + continue + } + if b.Len() > 0 { + b.WriteByte('\n') + } + b.WriteString(text) + } + return b.String() +} + +// decodeStringOrJoin accepts either a JSON string or a content-parts array +// (chat.completions multimodal) and returns a flat string. Multimodal parts +// are separated by newlines; non-text parts are skipped. +func decodeStringOrJoin(raw json.RawMessage) string { + if s, ok := tryDecodeString(raw); ok { + return s + } + return extractContentParts(raw) +} + +func tryDecodeString(raw json.RawMessage) (string, bool) { + if len(raw) == 0 { + return "", false + } + var s string + if err := json.Unmarshal(raw, &s); err == nil { + return s, true + } + return "", false +} + +// pickInt64 returns the first non-nil pointer's value. Used to prefer one +// naming convention while transparently falling back to another. +func pickInt64(preferred, fallback *int64) int64 { + if preferred != nil { + return *preferred + } + return derefInt64(fallback) +} + +func derefInt64(v *int64) int64 { + if v == nil { + return 0 + } + return *v +} + +func ptrDeref(b *bool) bool { + if b == nil { + return false + } + return *b +} + +func isEventStream(contentType string) bool { + return strings.Contains(strings.ToLower(contentType), "text/event-stream") +} + +func isJSON(contentType string) bool { + lower := strings.ToLower(contentType) + return strings.Contains(lower, "application/json") || strings.Contains(lower, "+json") +} diff --git a/proxy/internal/llm/openai_test.go b/proxy/internal/llm/openai_test.go new file mode 100644 index 000000000..7a5fca4fb --- /dev/null +++ b/proxy/internal/llm/openai_test.go @@ -0,0 +1,255 @@ +package llm + +import ( + "errors" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestOpenAIDetectFromURL(t *testing.T) { + p := OpenAIParser{} + + cases := map[string]bool{ + "/v1/chat/completions": true, + "/v1/completions": true, + "/v1/embeddings": true, + "/v1/responses": true, + "/API/V1/Chat/Completions": true, + "/upstream/v1/chat/completions?trace=1": true, + // Cloudflare AI Gateway puts its own /v1/{account}/{gateway} + // segment between the canonical /v1/ and the provider slug, + // so the /v1/chat/completions substring no longer appears + // adjacent in the path. The bare /chat/completions hint + // catches Cloudflare's OpenAI direct path + // (/v1/{account}/{gateway}/openai/chat/completions) and + // compat path (/v1/{account}/{gateway}/compat/chat/completions). + "/v1/{account}/{gateway}/openai/chat/completions": true, + "/v1/{account}/{gateway}/compat/chat/completions": true, + "/chat/completions": true, + "/v1/messages": false, + "/healthz": false, + "": false, + } + for path, want := range cases { + assert.Equal(t, want, p.DetectFromURL(path), "DetectFromURL(%q)", path) + } +} + +func TestOpenAIParseRequest(t *testing.T) { + p := OpenAIParser{} + + t.Run("stream true", func(t *testing.T) { + facts, err := p.ParseRequest([]byte(`{"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true}}`)) + require.NoError(t, err) + assert.Equal(t, "gpt-4o", facts.Model, "request model extracted") + assert.True(t, facts.Stream, "request marked as streaming") + }) + + t.Run("stream default", func(t *testing.T) { + facts, err := p.ParseRequest([]byte(`{"model":"gpt-4o-mini"}`)) + require.NoError(t, err) + assert.Equal(t, "gpt-4o-mini", facts.Model, "request model extracted") + assert.False(t, facts.Stream, "missing stream flag defaults to false") + }) + + t.Run("malformed", func(t *testing.T) { + _, err := p.ParseRequest([]byte(`{not json}`)) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrMalformedRequest), "sentinel error wrapped") + }) +} + +func TestOpenAIParseResponse(t *testing.T) { + p := OpenAIParser{} + + t.Run("happy fixture", func(t *testing.T) { + body, err := os.ReadFile(filepath.Join("fixtures", "openai_chat_completion.json")) + require.NoError(t, err, "fixture must be readable") + + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Equal(t, int64(123), usage.InputTokens, "prompt tokens become input") + assert.Equal(t, int64(45), usage.OutputTokens, "completion tokens become output") + assert.Equal(t, int64(168), usage.TotalTokens, "total_tokens carried through") + }) + + t.Run("total computed when missing", func(t *testing.T) { + body := []byte(`{"usage":{"prompt_tokens":10,"completion_tokens":5}}`) + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Equal(t, int64(15), usage.TotalTokens, "total computed from in+out") + }) + + t.Run("streaming rejected", func(t *testing.T) { + _, err := p.ParseResponse(200, "text/event-stream", []byte("")) + require.ErrorIs(t, err, ErrStreamingUnsupported, "SSE responses must use the scanner") + }) + + t.Run("non-200", func(t *testing.T) { + _, err := p.ParseResponse(500, "application/json", []byte(`{"error":"x"}`)) + require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 rejected as non-LLM") + }) + + t.Run("non-json content type", func(t *testing.T) { + _, err := p.ParseResponse(200, "text/plain", []byte(`{}`)) + require.ErrorIs(t, err, ErrNotLLMResponse, "text/plain treated as non-LLM") + }) + + t.Run("malformed body", func(t *testing.T) { + _, err := p.ParseResponse(200, "application/json", []byte(`{not json`)) + require.ErrorIs(t, err, ErrMalformedResponse, "bad JSON yields malformed error") + }) + + // Responses-API fixture: /v1/responses returns input_tokens/output_tokens + // (Anthropic-style) instead of prompt_tokens/completion_tokens. The parser + // must accept both. + t.Run("responses api fixture", func(t *testing.T) { + body, err := os.ReadFile(filepath.Join("fixtures", "openai_responses.json")) + require.NoError(t, err, "fixture must be readable") + + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Equal(t, int64(15), usage.InputTokens, "input_tokens should map directly") + assert.Equal(t, int64(414), usage.OutputTokens, "output_tokens should map directly") + assert.Equal(t, int64(429), usage.TotalTokens, "total_tokens carried through") + }) + + t.Run("responses api naming preferred over chat-completions when both present", func(t *testing.T) { + body := []byte(`{"usage":{"prompt_tokens":1,"completion_tokens":2,"input_tokens":15,"output_tokens":414,"total_tokens":429}}`) + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Equal(t, int64(15), usage.InputTokens, "responses-api names take precedence") + assert.Equal(t, int64(414), usage.OutputTokens, "responses-api names take precedence") + }) + + t.Run("chat-completions naming still works alone", func(t *testing.T) { + body := []byte(`{"usage":{"prompt_tokens":15,"completion_tokens":414,"total_tokens":429}}`) + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Equal(t, int64(15), usage.InputTokens, "prompt_tokens fallback") + assert.Equal(t, int64(414), usage.OutputTokens, "completion_tokens fallback") + }) + + // Cached-prompt accounting. cached_tokens is a SUBSET of + // prompt_tokens — input_tokens carries the full prompt count and + // the cached subset is reported separately so the cost meter can + // apply the discount rate to that portion. + t.Run("chat-completions cached_tokens subset surfaces", func(t *testing.T) { + body := []byte(`{"usage":{"prompt_tokens":1024,"completion_tokens":200,"total_tokens":1224,"prompt_tokens_details":{"cached_tokens":768}}}`) + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Equal(t, int64(1024), usage.InputTokens, "input remains the full prompt count — cached is a subset, not a separate bucket") + assert.Equal(t, int64(768), usage.CachedInputTokens, "cached_tokens must surface so cost meter can discount the cached subset") + assert.Zero(t, usage.CacheCreationTokens, "OpenAI has no cache_creation analogue") + }) + + t.Run("responses-api input_tokens_details.cached_tokens surfaces", func(t *testing.T) { + body := []byte(`{"usage":{"input_tokens":2048,"output_tokens":100,"total_tokens":2148,"input_tokens_details":{"cached_tokens":1500}}}`) + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Equal(t, int64(2048), usage.InputTokens) + assert.Equal(t, int64(1500), usage.CachedInputTokens, "Responses-API input_tokens_details.cached_tokens path must surface too") + }) + + t.Run("responses-api cached takes precedence over chat-completions when both present", func(t *testing.T) { + body := []byte(`{"usage":{"prompt_tokens":1,"input_tokens":2,"output_tokens":3,"prompt_tokens_details":{"cached_tokens":50},"input_tokens_details":{"cached_tokens":99}}}`) + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Equal(t, int64(99), usage.CachedInputTokens, "Responses-API field wins when both naming conventions are present") + }) + + t.Run("absent cached_tokens leaves cached counts at zero", func(t *testing.T) { + body := []byte(`{"usage":{"prompt_tokens":15,"completion_tokens":414,"total_tokens":429}}`) + usage, err := p.ParseResponse(200, "application/json", body) + require.NoError(t, err) + assert.Zero(t, usage.CachedInputTokens, "no prompt_tokens_details = no cached subset") + }) +} + +func TestOpenAIExtractPrompt_ChatCompletions(t *testing.T) { + body := []byte(`{"model":"gpt-4o-mini","messages":[{"role":"system","content":"be brief"},{"role":"user","content":"ping"}]}`) + got := OpenAIParser{}.ExtractPrompt(body) + require.NotEmpty(t, got, "messages array must extract") + require.Contains(t, got, "system: be brief", "system role and content surface") + require.Contains(t, got, "user: ping", "user role and content surface") +} + +func TestOpenAIExtractPrompt_ResponsesAPIStringInput(t *testing.T) { + body := []byte(`{"model":"gpt-5.4","input":"Hello there"}`) + got := OpenAIParser{}.ExtractPrompt(body) + require.Equal(t, "Hello there", got, "string input field should pass through") +} + +func TestOpenAIExtractPrompt_ResponsesAPIInputParts(t *testing.T) { + body := []byte(`{"model":"gpt-5.4","input":[{"type":"input_text","input_text":"first"},{"type":"input_text","input_text":"second"}]}`) + got := OpenAIParser{}.ExtractPrompt(body) + require.Contains(t, got, "first", "first content part surfaces") + require.Contains(t, got, "second", "second content part surfaces") +} + +// TestOpenAIExtractPrompt_ResponsesAPIMessageItems guards the live Codex +// shape: input is an array of message items whose text is nested under +// content[].text, not flat content parts. The old code fed the outer array +// to the content-part decoder and extracted nothing, so the stored prompt +// was empty. +func TestOpenAIExtractPrompt_ResponsesAPIMessageItems(t *testing.T) { + body := []byte(`{"model":"gpt-5.5","input":[` + + `{"type":"message","role":"developer","content":[{"type":"input_text","text":"system rules"}]},` + + `{"type":"message","role":"user","content":[{"type":"input_text","text":"hello there"}]},` + + `{"type":"reasoning","encrypted_content":"opaque","summary":[]},` + + `{"type":"message","role":"assistant","content":[{"type":"output_text","text":"prior reply"}]}` + + `]}`) + got := OpenAIParser{}.ExtractPrompt(body) + require.Contains(t, got, "system rules", "developer message content must surface") + require.Contains(t, got, "hello there", "user message content must surface") + require.Contains(t, got, "developer:", "role labels must prefix each message") + require.NotContains(t, got, "opaque", "reasoning items without text must be skipped") +} + +func TestOpenAIExtractPrompt_LegacyCompletion(t *testing.T) { + body := []byte(`{"model":"text-davinci-003","prompt":"once upon a time"}`) + got := OpenAIParser{}.ExtractPrompt(body) + require.Equal(t, "once upon a time", got, "string prompt field should pass through") +} + +func TestOpenAIExtractSessionID(t *testing.T) { + t.Run("codex client_metadata.session_id", func(t *testing.T) { + body := []byte(`{"model":"gpt-5.5","client_metadata":{"session_id":"019eeb72-ab7c-7cd2","thread_id":"t1"},"input":[]}`) + assert.Equal(t, "019eeb72-ab7c-7cd2", OpenAIParser{}.ExtractSessionID(body), "Codex session id must come from client_metadata.session_id") + }) + t.Run("plain chat has no session", func(t *testing.T) { + body := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`) + assert.Equal(t, "", OpenAIParser{}.ExtractSessionID(body), "plain chat.completions carries no session id") + }) + t.Run("non-JSON yields empty", func(t *testing.T) { + assert.Equal(t, "", OpenAIParser{}.ExtractSessionID([]byte("not json")), "malformed body must not error") + }) +} + +func TestOpenAIExtractCompletion_ChatCompletions(t *testing.T) { + body, err := os.ReadFile(filepath.Join("fixtures", "openai_chat_completion.json")) + require.NoError(t, err) + got := OpenAIParser{}.ExtractCompletion(200, "application/json", body) + require.NotEmpty(t, got, "fixture has assistant content") +} + +func TestOpenAIExtractCompletion_ResponsesAPI(t *testing.T) { + body, err := os.ReadFile(filepath.Join("fixtures", "openai_responses.json")) + require.NoError(t, err) + got := OpenAIParser{}.ExtractCompletion(200, "application/json", body) + require.NotEmpty(t, got, "responses-api fixture has output content") +} + +func TestOpenAIExtractCompletion_Streaming(t *testing.T) { + got := OpenAIParser{}.ExtractCompletion(200, "text/event-stream", []byte("")) + require.Empty(t, got, "streaming responses are skipped") +} + +func TestOpenAIExtractCompletion_NonOK(t *testing.T) { + got := OpenAIParser{}.ExtractCompletion(500, "application/json", []byte(`{"choices":[{"message":{"content":"x"}}]}`)) + require.Empty(t, got, "non-200 returns empty") +} diff --git a/proxy/internal/llm/parser.go b/proxy/internal/llm/parser.go new file mode 100644 index 000000000..81fa11f97 --- /dev/null +++ b/proxy/internal/llm/parser.go @@ -0,0 +1,112 @@ +// Package llm provides the shared LLM request and response parsing +// library consumed by proxy middleware. It is runtime agnostic: the same +// package is used by the native built-in executor now and will be reused +// by the WASM adapter later. +package llm + +// Provider identifies an LLM API provider. +type Provider int + +const ( + // ProviderUnknown signals that no parser matched the request. + ProviderUnknown Provider = 0 + // ProviderOpenAI identifies the OpenAI API surface. + ProviderOpenAI Provider = 1 + // ProviderAnthropic identifies the Anthropic Messages API surface. + ProviderAnthropic Provider = 2 + // ProviderBedrock identifies the AWS Bedrock runtime surface. + ProviderBedrock Provider = 3 +) + +// RequestFacts captures the subset of the LLM request body that the +// middleware annotates as metadata (model, streaming flag). Additional +// fields are added as parsers grow. +type RequestFacts struct { + Model string + Stream bool +} + +// Usage is the provider-agnostic token accounting emitted to metrics and +// access logs. Downstream consumers map InputTokens/OutputTokens to the +// plg.llm.* metadata allowlist entries. +// +// CachedInputTokens carries OpenAI's prompt_tokens_details.cached_tokens +// (a SUBSET of InputTokens) when the response is from OpenAI, or +// Anthropic's cache_read_input_tokens (ADDITIVE to InputTokens) when from +// Anthropic. The cost meter switches formula on KeyLLMProvider so the +// two shapes are billed correctly without double-counting. +// +// CacheCreationTokens carries Anthropic's cache_creation_input_tokens +// (ADDITIVE; not present in the OpenAI shape). +type Usage struct { + InputTokens int64 + OutputTokens int64 + TotalTokens int64 + CachedInputTokens int64 + CacheCreationTokens int64 +} + +// Parser is the per-provider interface implemented in this package. The +// dispatcher selects a parser by calling DetectFromURL against the incoming +// request path; ties break by registration order (see Parsers). +type Parser interface { + Provider() Provider + ProviderName() string + DetectFromURL(path string) bool + ParseRequest(body []byte) (RequestFacts, error) + ParseResponse(status int, contentType string, body []byte) (Usage, error) + // ExtractPrompt returns the user-facing prompt text from a request body. + // Different endpoint shapes (chat.completions, responses, messages) are + // handled by the per-provider implementation. Returns "" when no prompt + // can be extracted; never returns an error — extraction is best-effort + // because callers use the result for observability, not authorization. + ExtractPrompt(body []byte) string + // ExtractCompletion returns the assistant-facing completion text from a + // non-streaming response body. status and contentType match the + // ParseResponse arguments so implementations can fast-fail uniformly. + ExtractCompletion(status int, contentType string, body []byte) string + // ExtractSessionID returns a stable identifier that groups requests of + // the same conversation / coding session, read from the per-provider + // location clients populate (e.g. OpenAI Codex's client_metadata.session_id, + // Claude Code's metadata.user_id). Returns "" when the body carries no + // recognised session marker; extraction is best-effort and never errors. + ExtractSessionID(body []byte) string +} + +// Parsers returns the built-in parser set in a stable order. The order is +// deterministic so that DetectFromURL ties produce consistent routing. +func Parsers() []Parser { + return []Parser{ + OpenAIParser{}, + AnthropicParser{}, + BedrockParser{}, + } +} + +// DetectParser returns the first parser whose DetectFromURL matches the given +// request path. ok=false means no parser claimed the path. +func DetectParser(path string) (Parser, bool) { + for _, p := range Parsers() { + if p.DetectFromURL(path) { + return p, true + } + } + return nil, false +} + +// ParserByName returns the parser whose ProviderName matches id. Used by +// callers that already know which provider surface a request will hit +// (e.g. the agent-network middleware chain configured per synthesised +// service) so they can skip URL sniffing. ok=false when no parser is +// registered under that name. +func ParserByName(id string) (Parser, bool) { + if id == "" { + return nil, false + } + for _, p := range Parsers() { + if p.ProviderName() == id { + return p, true + } + } + return nil, false +} diff --git a/proxy/internal/llm/parser_test.go b/proxy/internal/llm/parser_test.go new file mode 100644 index 000000000..b3052ce68 --- /dev/null +++ b/proxy/internal/llm/parser_test.go @@ -0,0 +1,54 @@ +package llm + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParsers_ProviderNames(t *testing.T) { + parsers := Parsers() + require.Len(t, parsers, 3, "three built-in parsers expected") + + names := make([]string, 0, len(parsers)) + for _, p := range parsers { + names = append(names, p.ProviderName()) + } + assert.Contains(t, names, "openai", "OpenAI parser should be registered") + assert.Contains(t, names, "anthropic", "Anthropic parser should be registered") + assert.Contains(t, names, "bedrock", "Bedrock parser should be registered") +} + +func TestDetectParser(t *testing.T) { + cases := []struct { + name string + path string + expectedName string + expectOK bool + }{ + {"openai chat", "/v1/chat/completions", "openai", true}, + {"openai prefixed", "/api/v1/chat/completions", "openai", true}, + {"openai responses", "/v1/responses", "openai", true}, + {"anthropic messages", "/v1/messages", "anthropic", true}, + {"anthropic prefixed", "/proxy/v1/messages?query", "anthropic", true}, + {"unknown path", "/healthz", "", false}, + {"empty path", "", "", false}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + p, ok := DetectParser(tc.path) + require.Equal(t, tc.expectOK, ok, "detection success mismatch for %q", tc.path) + if ok { + assert.Equal(t, tc.expectedName, p.ProviderName(), "provider name mismatch") + } + }) + } +} + +func TestProviderValues(t *testing.T) { + assert.Equal(t, Provider(0), ProviderUnknown, "unknown provider is the zero value") + assert.Equal(t, ProviderOpenAI, OpenAIParser{}.Provider(), "OpenAI parser reports its provider enum") + assert.Equal(t, ProviderAnthropic, AnthropicParser{}.Provider(), "Anthropic parser reports its provider enum") +} diff --git a/proxy/internal/llm/pricing/defaults_coverage_test.go b/proxy/internal/llm/pricing/defaults_coverage_test.go new file mode 100644 index 000000000..be23682da --- /dev/null +++ b/proxy/internal/llm/pricing/defaults_coverage_test.go @@ -0,0 +1,65 @@ +package pricing + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestDefaultTable_FirstPartyModelCoverage guards the embedded defaults against +// silent drift/gaps: every metered first-party model the management catalog +// enumerates must resolve to a price, and a few rates that previously drifted +// are pinned to their LiteLLM-validated values. Keep this list in step with the +// catalog (management/server/agentnetwork/catalog) when adding models. +func TestDefaultTable_FirstPartyModelCoverage(t *testing.T) { + tbl := DefaultTable() + require.NotNil(t, tbl, "embedded default pricing table must load") + + mustPrice := map[string][]string{ + // openai parser covers openai_api, azure_openai_api, and mistral_api. + "openai": { + "gpt-5.5", "gpt-5.5-pro", "gpt-5.4", "gpt-5.4-mini", "gpt-5.4-nano", + "gpt-5.3-codex", "gpt-5.3-chat-latest", "o4-mini", + "gpt-4.1", "gpt-4.1-mini", "gpt-4.1-nano", "gpt-4o", "gpt-4o-mini", + "gpt-4-turbo", "gpt-3.5-turbo", "gpt-35-turbo", + "text-embedding-3-large", "text-embedding-3-small", + "mistral-large-latest", "mistral-medium-3-5", "codestral-2508", + "ministral-8b-latest", "mistral-embed", + }, + "anthropic": { + "claude-fable-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6", + "claude-opus-4-1", "claude-sonnet-4-6", "claude-sonnet-4-5", "claude-haiku-4-5", + }, + // bedrock keys are the normalized ids the request parser emits. + "bedrock": { + "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6", + "anthropic.claude-opus-4-1", "anthropic.claude-sonnet-4-6", "anthropic.claude-sonnet-4-5", + "anthropic.claude-haiku-4-5", "meta.llama3-3-70b-instruct", + "amazon.nova-pro", "amazon.nova-lite", "amazon.nova-micro", "amazon.nova-2-lite", + }, + } + for provider, models := range mustPrice { + for _, m := range models { + _, ok := tbl.Cost(provider, m, 1000, 1000, 0, 0) + assert.True(t, ok, "%s/%s must be priced in the embedded defaults", provider, m) + } + } + + // Pin per-direction rates independently (input-only then output-only) so a + // swap or skew of input<->output that preserves the combined total is still + // caught — these are rates that previously drifted or are easy to mis-enter. + in, ok := tbl.Cost("openai", "gpt-5.4", 1000, 0, 0, 0) + require.True(t, ok) + assert.InDelta(t, 0.0025, in, 1e-9, "gpt-5.4 input = 0.0025 per 1k") + out, ok := tbl.Cost("openai", "gpt-5.4", 0, 1000, 0, 0) + require.True(t, ok) + assert.InDelta(t, 0.015, out, 1e-9, "gpt-5.4 output = 0.015 per 1k") + + in, ok = tbl.Cost("bedrock", "anthropic.claude-sonnet-4-5", 1000, 0, 0, 0) + require.True(t, ok) + assert.InDelta(t, 0.003, in, 1e-9, "bedrock sonnet-4-5 input = 0.003 per 1k") + out, ok = tbl.Cost("bedrock", "anthropic.claude-sonnet-4-5", 0, 1000, 0, 0) + require.True(t, ok) + assert.InDelta(t, 0.015, out, 1e-9, "bedrock sonnet-4-5 output = 0.015 per 1k") +} diff --git a/proxy/internal/llm/pricing/defaults_pricing.yaml b/proxy/internal/llm/pricing/defaults_pricing.yaml new file mode 100644 index 000000000..cd5c64fbf --- /dev/null +++ b/proxy/internal/llm/pricing/defaults_pricing.yaml @@ -0,0 +1,264 @@ +# Embedded default pricing for llm_observability. Compiled into the proxy +# binary via go:embed in pricing.go; cost annotation works out of the box +# without any operator action. +# +# Operators override entries by dropping a pricing.yaml into --plugin-data-dir +# (or whichever basename is given via params.pricing_path). The override file +# only needs entries the operator wants to change; missing entries fall +# through to these defaults. +# +# Values are USD per 1_000 tokens. Public list prices drift; ship a fresh +# binary or override individual entries via the override file as needed. +# +# Optional cache fields: +# cached_input_per_1k OpenAI: rate for prompt_tokens_details.cached_tokens +# (a SUBSET of prompt_tokens). Typically 0.5x input. +# Absent → cached portion bills at input_per_1k. +# cache_read_per_1k Anthropic: rate for cache_read_input_tokens +# (ADDITIVE to input_tokens). Typically 0.1x input. +# Absent → cache reads bill at input_per_1k. +# cache_creation_per_1k Anthropic: rate for cache_creation_input_tokens +# (ADDITIVE to input_tokens). Typically 1.25x input. +# Absent → cache writes bill at input_per_1k. + +openai: + # OpenAI + OpenAI-compatible providers (openai_api, azure_openai_api, + # mistral_api, and the openai-parser gateways) all emit llm.provider="openai", + # so their models are priced here. Kept in sync with the management catalog; + # rates cross-checked against LiteLLM model_prices_and_context_window.json. + + # GPT-5.x family — cache reads 10% of input (0.1x). + gpt-5.5: + input_per_1k: 0.005 + output_per_1k: 0.03 + cached_input_per_1k: 0.0005 + gpt-5.5-pro: + input_per_1k: 0.03 + output_per_1k: 0.18 + cached_input_per_1k: 0.003 + gpt-5.4: + input_per_1k: 0.0025 + output_per_1k: 0.015 + cached_input_per_1k: 0.00025 + gpt-5.4-pro: + input_per_1k: 0.03 + output_per_1k: 0.18 + cached_input_per_1k: 0.003 + gpt-5.4-mini: + input_per_1k: 0.00075 + output_per_1k: 0.0045 + cached_input_per_1k: 0.000075 + gpt-5.4-nano: + input_per_1k: 0.0002 + output_per_1k: 0.00125 + cached_input_per_1k: 0.00002 + gpt-5.3-codex: + input_per_1k: 0.00175 + output_per_1k: 0.014 + cached_input_per_1k: 0.000175 + gpt-5.3-chat-latest: + input_per_1k: 0.00175 + output_per_1k: 0.014 + cached_input_per_1k: 0.000175 + # GPT-5 (2025) family — kept for gateway requests using the unsuffixed ids. + gpt-5: + input_per_1k: 0.00125 + output_per_1k: 0.01 + cached_input_per_1k: 0.000125 + gpt-5-mini: + input_per_1k: 0.00025 + output_per_1k: 0.002 + cached_input_per_1k: 0.000025 + gpt-5-nano: + input_per_1k: 0.00005 + output_per_1k: 0.0004 + cached_input_per_1k: 0.000005 + o4-mini: + input_per_1k: 0.0011 + output_per_1k: 0.0044 + cached_input_per_1k: 0.000275 + # GPT-4.1 family — cache reads 25% of input. + gpt-4.1: + input_per_1k: 0.002 + output_per_1k: 0.008 + cached_input_per_1k: 0.0005 + gpt-4.1-mini: + input_per_1k: 0.0004 + output_per_1k: 0.0016 + cached_input_per_1k: 0.0001 + gpt-4.1-nano: + input_per_1k: 0.0001 + output_per_1k: 0.0004 + cached_input_per_1k: 0.000025 + # GPT-4o family — cache reads 50% of input (0.5x). + gpt-4o: + input_per_1k: 0.0025 + output_per_1k: 0.01 + cached_input_per_1k: 0.00125 + gpt-4o-mini: + input_per_1k: 0.00015 + output_per_1k: 0.0006 + cached_input_per_1k: 0.000075 + # Older GPT — no prompt caching. + gpt-4-turbo: + input_per_1k: 0.01 + output_per_1k: 0.03 + gpt-3.5-turbo: + input_per_1k: 0.0005 + output_per_1k: 0.0015 + gpt-35-turbo: # Azure deployment alias of gpt-3.5-turbo + input_per_1k: 0.0005 + output_per_1k: 0.0015 + # Embeddings — no caching, no output tokens. + text-embedding-3-large: + input_per_1k: 0.00013 + output_per_1k: 0 + text-embedding-3-small: + input_per_1k: 0.00002 + output_per_1k: 0 + + # Mistral (mistral_api) — routed via the openai parser; no prompt caching. + mistral-large-latest: + input_per_1k: 0.0005 + output_per_1k: 0.0015 + mistral-medium-latest: + input_per_1k: 0.0004 + output_per_1k: 0.002 + mistral-medium-3-5: + input_per_1k: 0.0015 + output_per_1k: 0.0075 + mistral-small-latest: + input_per_1k: 0.00006 + output_per_1k: 0.00018 + magistral-medium-latest: + input_per_1k: 0.002 + output_per_1k: 0.005 + magistral-small-latest: + input_per_1k: 0.0005 + output_per_1k: 0.0015 + devstral-medium-latest: + input_per_1k: 0.0004 + output_per_1k: 0.002 + devstral-small-latest: + input_per_1k: 0.0001 + output_per_1k: 0.0003 + codestral-2508: + input_per_1k: 0.0003 + output_per_1k: 0.0009 + codestral-latest: + input_per_1k: 0.001 + output_per_1k: 0.003 + ministral-3-14b-2512: + input_per_1k: 0.0002 + output_per_1k: 0.0002 + ministral-8b-latest: + input_per_1k: 0.00015 + output_per_1k: 0.00015 + ministral-3-3b-2512: + input_per_1k: 0.0001 + output_per_1k: 0.0001 + mistral-embed: + input_per_1k: 0.0001 + output_per_1k: 0 + +anthropic: + # Claude 4.x family — cache reads ≈10% of input, cache writes ≈125% of input. + # Pricing source: Anthropic's current published rates per million tokens, + # divided by 1000 for the per-1k figures stored here. + claude-fable-5: + input_per_1k: 0.010 + output_per_1k: 0.050 + cache_read_per_1k: 0.001 + cache_creation_per_1k: 0.0125 + claude-opus-4-8: + input_per_1k: 0.005 + output_per_1k: 0.025 + cache_read_per_1k: 0.0005 + cache_creation_per_1k: 0.00625 + claude-opus-4-7: + input_per_1k: 0.005 + output_per_1k: 0.025 + cache_read_per_1k: 0.0005 + cache_creation_per_1k: 0.00625 + claude-opus-4-6: + input_per_1k: 0.005 + output_per_1k: 0.025 + cache_read_per_1k: 0.0005 + cache_creation_per_1k: 0.00625 + claude-opus-4-1: + input_per_1k: 0.015 + output_per_1k: 0.075 + cache_read_per_1k: 0.0015 + cache_creation_per_1k: 0.01875 + claude-sonnet-4-6: + input_per_1k: 0.003 + output_per_1k: 0.015 + cache_read_per_1k: 0.0003 + cache_creation_per_1k: 0.00375 + claude-sonnet-4-5: + input_per_1k: 0.003 + output_per_1k: 0.015 + cache_read_per_1k: 0.0003 + cache_creation_per_1k: 0.00375 + claude-haiku-4-5: + input_per_1k: 0.001 + output_per_1k: 0.005 + cache_read_per_1k: 0.0001 + cache_creation_per_1k: 0.00125 + +bedrock: + # AWS Bedrock model ids, normalised by the request parser (cross-region + # inference-profile prefix + version/throughput suffix stripped), e.g. + # eu.anthropic.claude-sonnet-4-5-20250929-v1:0 -> anthropic.claude-sonnet-4-5. + # Anthropic-on-Bedrock keeps the additive cache buckets (read ≈0.1x input, + # write ≈1.25x input); Nova / Llama report no cache, so cost is input+output. + anthropic.claude-opus-4-8: + input_per_1k: 0.005 + output_per_1k: 0.025 + cache_read_per_1k: 0.0005 + cache_creation_per_1k: 0.00625 + anthropic.claude-opus-4-7: + input_per_1k: 0.005 + output_per_1k: 0.025 + cache_read_per_1k: 0.0005 + cache_creation_per_1k: 0.00625 + anthropic.claude-opus-4-6: + input_per_1k: 0.005 + output_per_1k: 0.025 + cache_read_per_1k: 0.0005 + cache_creation_per_1k: 0.00625 + anthropic.claude-opus-4-1: + input_per_1k: 0.015 + output_per_1k: 0.075 + cache_read_per_1k: 0.0015 + cache_creation_per_1k: 0.01875 + anthropic.claude-sonnet-4-6: + input_per_1k: 0.003 + output_per_1k: 0.015 + cache_read_per_1k: 0.0003 + cache_creation_per_1k: 0.00375 + anthropic.claude-sonnet-4-5: + input_per_1k: 0.003 + output_per_1k: 0.015 + cache_read_per_1k: 0.0003 + cache_creation_per_1k: 0.00375 + anthropic.claude-haiku-4-5: + input_per_1k: 0.001 + output_per_1k: 0.005 + cache_read_per_1k: 0.0001 + cache_creation_per_1k: 0.00125 + meta.llama3-3-70b-instruct: + input_per_1k: 0.00072 + output_per_1k: 0.00072 + amazon.nova-2-lite: + input_per_1k: 0.0003 + output_per_1k: 0.0025 + amazon.nova-pro: + input_per_1k: 0.0008 + output_per_1k: 0.0032 + amazon.nova-lite: + input_per_1k: 0.00006 + output_per_1k: 0.00024 + amazon.nova-micro: + input_per_1k: 0.000035 + output_per_1k: 0.00014 diff --git a/proxy/internal/llm/pricing/pricing.go b/proxy/internal/llm/pricing/pricing.go new file mode 100644 index 000000000..09afec5ff --- /dev/null +++ b/proxy/internal/llm/pricing/pricing.go @@ -0,0 +1,449 @@ +// Package pricing implements the embedded-default + override pricing table +// shared by middleware that converts LLM token usage into a USD cost +// estimate. The table is hot-reloadable from a basename under the proxy +// data directory; missing override files keep the embedded defaults so +// cost annotation works without operator action. +package pricing + +import ( + "bytes" + "context" + _ "embed" + "errors" + "fmt" + "io" + "io/fs" + "math" + "path/filepath" + "regexp" + "strings" + "sync" + "sync/atomic" + "time" + + log "github.com/sirupsen/logrus" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + "gopkg.in/yaml.v3" +) + +//go:embed defaults_pricing.yaml +var defaultPricingYAML []byte + +var ( + defaultTableOnce sync.Once + defaultTablePtr *Table +) + +// DefaultTable returns the pricing table embedded in the binary. The result +// is parsed once and shared; callers must not mutate the returned value. +// Cost annotation works without any operator action because every loader +// starts with this table. +func DefaultTable() *Table { + defaultTableOnce.Do(func() { + t, err := parsePricingBytes(defaultPricingYAML) + if err != nil { + panic(fmt.Sprintf("llmobs: embedded default pricing failed to parse: %v", err)) + } + defaultTablePtr = t + }) + return defaultTablePtr +} + +// mergeOver returns a new Table containing every entry from base, with any +// matching entry from overlay replacing the base value. Either argument may +// be nil. Result is a fresh allocation so callers can mutate / Store safely. +func mergeOver(base, overlay *Table) *Table { + if overlay == nil || len(overlay.entries) == 0 { + return base + } + if base == nil || len(base.entries) == 0 { + return overlay + } + out := make(map[string]map[string]Entry, len(base.entries)) + for provider, models := range base.entries { + inner := make(map[string]Entry, len(models)) + for model, e := range models { + inner[model] = e + } + out[provider] = inner + } + for provider, models := range overlay.entries { + inner, ok := out[provider] + if !ok { + inner = make(map[string]Entry, len(models)) + out[provider] = inner + } + for model, e := range models { + inner[model] = e + } + } + return &Table{entries: out} +} + +// Entry is a single model's input and output pricing, expressed in USD per +// 1000 tokens. +// +// CachedInputPer1K applies to OpenAI's cached prompt tokens, which are a +// subset of input_tokens — when set, the cached portion is billed at this +// rate and the non-cached remainder at InputPer1K. Zero means "no discount +// configured", and cached tokens are billed at InputPer1K (matches current +// behaviour where cached counts weren't extracted at all). +// +// CacheReadPer1K and CacheCreationPer1K apply to Anthropic's two prompt- +// cache fields, which are additive to input_tokens: cache_read is the +// cheaper read-from-cache rate, cache_creation is the more expensive +// write-to-cache rate. Zero means "no rate configured" and the +// corresponding token bucket is billed at InputPer1K. This is more +// accurate than today's behaviour, where Anthropic's cache tokens are +// ignored and not charged at all. +type Entry struct { + InputPer1K float64 + OutputPer1K float64 + CachedInputPer1K float64 + CacheReadPer1K float64 + CacheCreationPer1K float64 +} + +// Table is a provider-to-model pricing lookup. Instances are immutable once +// built and are swapped atomically by Loader. +type Table struct { + entries map[string]map[string]Entry +} + +// Cost returns the estimated USD cost for the given token counts. ok is +// false when the provider or model is not present in the table; the caller +// can still emit token metrics with a model=unknown label. +// +// Provider-shape semantics for cached / cache-creation counts: +// +// - OpenAI: cachedInput is a SUBSET of inTokens. The cached portion is +// billed at CachedInputPer1K (or InputPer1K when no override), and the +// non-cached remainder of inTokens at InputPer1K. cacheCreation is +// ignored (OpenAI has no analogue). +// - Anthropic: cachedInput (cache_read) and cacheCreation are ADDITIVE to +// inTokens. The three buckets are billed at CacheReadPer1K, +// CacheCreationPer1K, and InputPer1K respectively, each falling back +// to InputPer1K when the corresponding rate is zero. +// - Other providers: cached and cacheCreation are ignored; cost is +// inTokens*InputPer1K + outTokens*OutputPer1K. +func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (float64, bool) { + // Clamp negatives to zero before any pricing math so a malformed + // upstream count can never produce a negative cost. + if inTokens < 0 { + inTokens = 0 + } + if outTokens < 0 { + outTokens = 0 + } + if cachedInput < 0 { + cachedInput = 0 + } + if cacheCreation < 0 { + cacheCreation = 0 + } + if t == nil { + return 0, false + } + byModel, ok := t.entries[provider] + if !ok { + return 0, false + } + entry, ok := byModel[model] + if !ok { + return 0, false + } + output := (float64(outTokens) / 1000.0) * entry.OutputPer1K + switch provider { + case "openai": + // cachedInput is a subset of inTokens; clamp so a malformed + // upstream (cached > total) can't produce a negative remainder. + clamped := cachedInput + if clamped > inTokens { + clamped = inTokens + } + cachedRate := entry.CachedInputPer1K + if cachedRate <= 0 { + cachedRate = entry.InputPer1K + } + nonCached := float64(inTokens-clamped) / 1000.0 * entry.InputPer1K + cached := float64(clamped) / 1000.0 * cachedRate + return nonCached + cached + output, true + case "anthropic", "bedrock": + // Bedrock-Anthropic returns the same additive cache buckets as + // first-party Anthropic; non-Anthropic Bedrock models simply report + // zero cache tokens, so this formula degrades to input + output. + readRate := entry.CacheReadPer1K + if readRate <= 0 { + readRate = entry.InputPer1K + } + createRate := entry.CacheCreationPer1K + if createRate <= 0 { + createRate = entry.InputPer1K + } + input := float64(inTokens) / 1000.0 * entry.InputPer1K + read := float64(cachedInput) / 1000.0 * readRate + create := float64(cacheCreation) / 1000.0 * createRate + return input + read + create + output, true + default: + input := float64(inTokens) / 1000.0 * entry.InputPer1K + return input + output, true + } +} + +// Has reports whether the provider/model pair is present in the table. +func (t *Table) Has(provider, model string) bool { + if t == nil { + return false + } + byModel, ok := t.entries[provider] + if !ok { + return false + } + _, ok = byModel[model] + return ok +} + +// pricingFile mirrors the on-disk YAML schema. Keys are provider names; the +// nested map keys are model names. +type pricingFile map[string]map[string]struct { + InputPer1K float64 `yaml:"input_per_1k"` + OutputPer1K float64 `yaml:"output_per_1k"` + CachedInputPer1K float64 `yaml:"cached_input_per_1k"` + CacheReadPer1K float64 `yaml:"cache_read_per_1k"` + CacheCreationPer1K float64 `yaml:"cache_creation_per_1k"` +} + +const ( + // ReloadInterval is the mtime-poll cadence for the background reloader. + ReloadInterval = 30 * time.Second + + // errorBackoff bounds how often the loader logs a repeated parse error. + errorBackoff = 5 * time.Minute +) + +var basenameRegex = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`) + +// Loader is a confined, hot-reloadable pricing table reader. Construction +// must succeed against the target file; subsequent reload failures keep the +// previously-loaded table so callers never observe a blank price list. +type Loader struct { + baseDir string + fullPath string + pluginID string + table atomic.Pointer[Table] + mtime atomic.Int64 + failures metric.Int64Counter + interval time.Duration +} + +// NewLoader returns a pricing loader that overlays an optional file-based +// table on top of the embedded defaults. Missing override file, baseDir, or +// relPath is not an error: the loader keeps the embedded defaults so cost +// metadata is still emitted for known models. +// +// Errors: +// - bad basename, traversal segment, or absolute relPath are rejected so a +// misconfigured target surfaces immediately. +// - permission errors and YAML parse errors keep the defaults but log a +// warning; cost annotation does not silently break. +// +// failures is optional; pass nil in tests that do not care about +// reload-failure telemetry. +func NewLoader(baseDir, relPath, pluginID string, failures metric.Int64Counter) (*Loader, error) { + defaults := DefaultTable() + l := &Loader{ + baseDir: baseDir, + pluginID: pluginID, + failures: failures, + } + l.table.Store(defaults) + + if strings.TrimSpace(baseDir) == "" || strings.TrimSpace(relPath) == "" { + return l, nil + } + + full, err := resolveMiddlewareDataPath(baseDir, relPath) + if err != nil { + return nil, err + } + l.fullPath = full + + overlay, mtime, err := loadPricing(full) + if err != nil { + if errors.Is(err, fs.ErrNotExist) { + // Override file is optional. Defaults already stored. + return l, nil + } + // Symlink rejection, oversize file, parse failure, permission errors + // — surface so a misconfigured operator sees the problem instead of + // silently running with stale defaults. + return nil, fmt.Errorf("load pricing %s: %w", full, err) + } + l.table.Store(mergeOver(defaults, overlay)) + l.mtime.Store(mtime.UnixNano()) + return l, nil +} + +// Get returns the current pricing table. The returned pointer is immutable; +// callers must not mutate its contents. +func (l *Loader) Get() *Table { + if l == nil { + return nil + } + return l.table.Load() +} + +// WatchesFile reports whether this loader is bound to an override file on +// disk. False for defaults-only loaders (no operator override given). +// Callers use this to decide whether to spawn the mtime-poll goroutine. +func (l *Loader) WatchesFile() bool { + if l == nil { + return false + } + return l.fullPath != "" +} + +// SetReloadInterval overrides the mtime-poll cadence used by Reload. Calls +// after Reload has started have no effect on the running loop. Intended for +// tests; production code uses the default ReloadInterval. +func (l *Loader) SetReloadInterval(d time.Duration) { + if l == nil || d <= 0 { + return + } + l.interval = d +} + +// Reload runs a polling loop that checks the pricing file mtime every +// ReloadInterval (or the value passed to SetReloadInterval). Returns when +// ctx is cancelled. +func (l *Loader) Reload(ctx context.Context) { + if l == nil { + return + } + interval := l.interval + if interval <= 0 { + interval = ReloadInterval + } + t := time.NewTicker(interval) + defer t.Stop() + + var lastErrAt time.Time + for { + select { + case <-ctx.Done(): + return + case <-t.C: + if err := l.reload(); err != nil { + if l.failures != nil { + l.failures.Add(ctx, 1, metric.WithAttributes( + attribute.String("plugin", l.pluginID), + )) + } + now := time.Now() + if now.Sub(lastErrAt) >= errorBackoff { + log.Warnf("llmobs: pricing reload failed for %s: %v", l.fullPath, err) + lastErrAt = now + } + } + } + } +} + +// reload performs a single-shot mtime check and reload. The reloaded +// override file is merged on top of the embedded defaults; missing override +// (e.g. operator deleted the file) is not an error and reverts to defaults. +func (l *Loader) reload() error { + if l.fullPath == "" { + // Defaults-only loader; nothing on disk to reload. + return nil + } + mtime, err := statMtime(l.fullPath) + if err != nil { + if errors.Is(err, fs.ErrNotExist) { + // File was removed since startup. Drop back to defaults and + // reset mtime so a future re-creation triggers a reload. + l.table.Store(DefaultTable()) + l.mtime.Store(0) + return nil + } + return err + } + if mtime.UnixNano() == l.mtime.Load() { + return nil + } + + overlay, newMtime, err := loadPricing(l.fullPath) + if err != nil { + return err + } + l.table.Store(mergeOver(DefaultTable(), overlay)) + l.mtime.Store(newMtime.UnixNano()) + return nil +} + +// resolveMiddlewareDataPath validates relPath is a safe basename and resolves +// it under baseDir. An additional cleaned-prefix check guards against +// CVE-style edge cases where Join is used with trailing path segments. +func resolveMiddlewareDataPath(baseDir, relPath string) (string, error) { + if strings.TrimSpace(baseDir) == "" { + return "", errors.New("middleware-data-dir is not configured") + } + if relPath == "" { + return "", errors.New("pricing path is empty") + } + if !basenameRegex.MatchString(relPath) { + return "", fmt.Errorf("pricing path %q is not a safe basename", relPath) + } + if filepath.IsAbs(relPath) { + return "", fmt.Errorf("pricing path %q must be a basename, not absolute", relPath) + } + + cleanBase, err := filepath.Abs(filepath.Clean(baseDir)) + if err != nil { + return "", fmt.Errorf("resolve middleware-data-dir: %w", err) + } + full := filepath.Join(cleanBase, relPath) + cleanedFull := filepath.Clean(full) + if !strings.HasPrefix(cleanedFull, cleanBase+string(filepath.Separator)) && cleanedFull != cleanBase { + return "", fmt.Errorf("pricing path %q escapes middleware-data-dir", relPath) + } + return cleanedFull, nil +} + +func parsePricingBytes(data []byte) (*Table, error) { + dec := yaml.NewDecoder(bytes.NewReader(data)) + dec.KnownFields(true) + + var raw pricingFile + if err := dec.Decode(&raw); err != nil && !errors.Is(err, io.EOF) { + return nil, fmt.Errorf("decode pricing yaml: %w", err) + } + + out := make(map[string]map[string]Entry, len(raw)) + for provider, models := range raw { + inner := make(map[string]Entry, len(models)) + for model, entry := range models { + for field, v := range map[string]float64{ + "input_per_1k": entry.InputPer1K, + "output_per_1k": entry.OutputPer1K, + "cached_input_per_1k": entry.CachedInputPer1K, + "cache_read_per_1k": entry.CacheReadPer1K, + "cache_creation_per_1k": entry.CacheCreationPer1K, + } { + if v < 0 || math.IsNaN(v) || math.IsInf(v, 0) { + return nil, fmt.Errorf("pricing %s/%s: %s must be a finite, non-negative rate, got %v", provider, model, field, v) + } + } + inner[model] = Entry{ + InputPer1K: entry.InputPer1K, + OutputPer1K: entry.OutputPer1K, + CachedInputPer1K: entry.CachedInputPer1K, + CacheReadPer1K: entry.CacheReadPer1K, + CacheCreationPer1K: entry.CacheCreationPer1K, + } + } + out[provider] = inner + } + return &Table{entries: out}, nil +} diff --git a/proxy/internal/llm/pricing/pricing_other.go b/proxy/internal/llm/pricing/pricing_other.go new file mode 100644 index 000000000..e65fffff1 --- /dev/null +++ b/proxy/internal/llm/pricing/pricing_other.go @@ -0,0 +1,20 @@ +//go:build !unix + +package pricing + +import ( + "fmt" + "time" +) + +// loadPricing is unavailable on non-Unix platforms because O_NOFOLLOW and +// fstat-from-FD are required to honour the spec's symlink-safety rules. The +// proxy is only deployed on Linux today; a Windows port would need an +// equivalent path-as-handle implementation. +func loadPricing(path string) (*Table, time.Time, error) { + return nil, time.Time{}, fmt.Errorf("llmobs pricing loader is not supported on this platform: %s", path) +} + +func statMtime(path string) (time.Time, error) { + return time.Time{}, fmt.Errorf("llmobs pricing loader is not supported on this platform: %s", path) +} diff --git a/proxy/internal/llm/pricing/pricing_test.go b/proxy/internal/llm/pricing/pricing_test.go new file mode 100644 index 000000000..7ac2a85dc --- /dev/null +++ b/proxy/internal/llm/pricing/pricing_test.go @@ -0,0 +1,432 @@ +//go:build unix + +package pricing + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func copyFixture(t *testing.T, src, dst string) { + t.Helper() + data, err := os.ReadFile(src) + require.NoError(t, err, "read source fixture") + require.NoError(t, os.WriteFile(dst, data, 0o600), "write target fixture") +} + +func TestNewLoader_HappyPath(t *testing.T) { + base := t.TempDir() + copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml")) + + l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil) + require.NoError(t, err, "NewLoader must succeed with a valid fixture") + table := l.Get() + require.NotNil(t, table, "table populated after load") + + cost, ok := table.Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0) + require.True(t, ok, "known provider/model resolves") + assert.InDelta(t, 0.00075, cost, 1e-9, "cost = 0.00015 + 0.0006 per 1k tokens") + + cost, ok = table.Cost("openai", "gpt-4o", 2000, 1000, 0, 0) + require.True(t, ok, "second known model resolves") + assert.InDelta(t, 0.015, cost, 1e-9, "cost for gpt-4o: 2*0.0025 + 1*0.01") + + cost, ok = table.Cost("anthropic", "claude-sonnet-4-5", 1000, 1000, 0, 0) + require.True(t, ok, "anthropic model resolves") + assert.InDelta(t, 0.018, cost, 1e-9, "cost for claude-sonnet-4-5: 0.003 + 0.015") +} + +// TestCost_OpenAICachedSubsetDiscount proves OpenAI's cached input +// tokens are billed at the configured cached_input_per_1k rate while +// the non-cached remainder of input_tokens is billed at the regular +// rate. Critical because OpenAI returns cached_tokens as a SUBSET of +// prompt_tokens — naïvely charging the cached count on top of +// prompt_tokens would double-bill that portion. +func TestCost_OpenAICachedSubsetDiscount(t *testing.T) { + tbl := &Table{entries: map[string]map[string]Entry{ + "openai": {"gpt-4o": { + InputPer1K: 0.0025, // 0.0025 USD per 1k input tokens + OutputPer1K: 0.01, + CachedInputPer1K: 0.00125, // 0.5x discount on cached + }}, + }} + // 1000 prompt tokens, 750 of which were cached. 250 non-cached + // at regular rate, 750 cached at the discount rate, 500 output. + cost, ok := tbl.Cost("openai", "gpt-4o", 1000, 500, 750, 0) + require.True(t, ok, "known model resolves") + want := (250.0/1000.0)*0.0025 + (750.0/1000.0)*0.00125 + (500.0/1000.0)*0.01 + assert.InDelta(t, want, cost, 1e-12, + "cached subset must bill at the discount rate; non-cached remainder at regular rate") +} + +// TestCost_OpenAICachedFallsBackToInputRate covers the operator +// opt-in contract: when CachedInputPer1K is unset (zero), cached +// tokens bill at the regular input rate. This matches today's +// behaviour (cached counts weren't extracted at all so they +// implicitly billed at the input rate via prompt_tokens). +func TestCost_OpenAICachedFallsBackToInputRate(t *testing.T) { + tbl := &Table{entries: map[string]map[string]Entry{ + "openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01}}, + }} + cost, ok := tbl.Cost("openai", "gpt-4o", 1000, 500, 750, 0) + require.True(t, ok) + want := 0.0025 + (500.0/1000.0)*0.01 + assert.InDelta(t, want, cost, 1e-12, + "absent cached_input_per_1k rate must fall back to input_per_1k — same as pre-feature behaviour") +} + +// TestCost_OpenAIClampsCachedToInputCount is the defensive guard +// against malformed upstream responses that report cached_tokens > +// prompt_tokens. We clamp so the formula never produces a negative +// "non-cached remainder" multiplied by the input rate. +func TestCost_OpenAIClampsCachedToInputCount(t *testing.T) { + tbl := &Table{entries: map[string]map[string]Entry{ + "openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01, CachedInputPer1K: 0.00125}}, + }} + cost, ok := tbl.Cost("openai", "gpt-4o", 100, 0, 9999, 0) + require.True(t, ok) + // All 100 cached, 0 non-cached. Output is 0. + want := (100.0 / 1000.0) * 0.00125 + assert.InDelta(t, want, cost, 1e-12, + "cached count > input count must clamp to input — never bill negative non-cached tokens") +} + +// TestCost_AnthropicCacheReadAndCreationAreAdditive proves the +// Anthropic shape: cache_read and cache_creation tokens are +// ADDITIVE to input_tokens (not subset), each billed at its own +// configured rate. The two rates pull in opposite directions — +// cache_read is the cheaper read-from-cache rate (≈0.1× input), +// cache_creation is the more expensive write-to-cache rate +// (≈1.25× input). +func TestCost_AnthropicCacheReadAndCreationAreAdditive(t *testing.T) { + tbl := &Table{entries: map[string]map[string]Entry{ + "anthropic": {"claude-sonnet": { + InputPer1K: 0.003, + OutputPer1K: 0.015, + CacheReadPer1K: 0.0003, // 0.1x of input + CacheCreationPer1K: 0.00375, // 1.25x of input + }}, + }} + // 256 regular input + 768 cache_read + 512 cache_creation + + // 200 output. Each input bucket bills at its own rate. + cost, ok := tbl.Cost("anthropic", "claude-sonnet", 256, 200, 768, 512) + require.True(t, ok, "known model resolves") + want := (256.0/1000.0)*0.003 + + (768.0/1000.0)*0.0003 + + (512.0/1000.0)*0.00375 + + (200.0/1000.0)*0.015 + assert.InDelta(t, want, cost, 1e-12, + "each Anthropic input bucket must bill at its own configured rate") +} + +// TestCost_AnthropicCacheRatesFallBackToInput covers the no-opt-in +// path: when neither CacheReadPer1K nor CacheCreationPer1K is set, +// cache tokens bill at the regular input rate. This is more +// accurate than today's behaviour (cache tokens ignored entirely) +// without requiring operators to opt in via YAML. +func TestCost_AnthropicCacheRatesFallBackToInput(t *testing.T) { + tbl := &Table{entries: map[string]map[string]Entry{ + "anthropic": {"claude-sonnet": {InputPer1K: 0.003, OutputPer1K: 0.015}}, + }} + cost, ok := tbl.Cost("anthropic", "claude-sonnet", 256, 200, 768, 512) + require.True(t, ok) + // Without overrides: every input bucket at input_per_1k. + want := ((256.0+768.0+512.0)/1000.0)*0.003 + (200.0/1000.0)*0.015 + assert.InDelta(t, want, cost, 1e-12, + "absent cache rates must fall back to input_per_1k — Anthropic cache tokens were ignored before this change, billing at input rate is more accurate as a default") +} + +func TestNewLoader_UnknownModel(t *testing.T) { + base := t.TempDir() + copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml")) + + l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil) + require.NoError(t, err) + + _, ok := l.Get().Cost("openai", "fantasy-model", 10, 10, 0, 0) + assert.False(t, ok, "unknown model returns ok=false") + + _, ok = l.Get().Cost("cohere", "anything", 10, 10, 0, 0) + assert.False(t, ok, "unknown provider returns ok=false") +} + +func TestNewLoader_InvalidYAMLRejected(t *testing.T) { + base := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(base, "pricing.yaml"), []byte("\t- this is not: valid: yaml: :["), 0o600)) + + _, err := NewLoader(base, "pricing.yaml", "llm_observability", nil) + require.Error(t, err, "invalid YAML must surface as construction error") +} + +func TestLoader_ReloadKeepsPreviousOnParseError(t *testing.T) { + base := t.TempDir() + target := filepath.Join(base, "pricing.yaml") + copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target) + + l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil) + require.NoError(t, err) + require.NotNil(t, l.Get(), "initial table populated") + + // Overwrite with content that violates the strict schema (extra field) + // plus a bumped mtime to trigger reload. + require.NoError(t, os.WriteFile(target, []byte("openai:\n gpt-4o:\n input_per_1k: 1.0\n output_per_1k: 2.0\n bogus_field: nope\n"), 0o600)) + future := time.Now().Add(time.Hour) + require.NoError(t, os.Chtimes(target, future, future)) + + err = l.reload() + require.Error(t, err, "parse error surfaced by reload()") + + cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0) + require.True(t, ok, "previous table still available after parse failure") + assert.InDelta(t, 0.00075, cost, 1e-9, "previous cost preserved") +} + +func TestLoader_ReloadNoChangeIsNoOp(t *testing.T) { + base := t.TempDir() + target := filepath.Join(base, "pricing.yaml") + copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target) + + l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil) + require.NoError(t, err) + ptrBefore := l.Get() + + require.NoError(t, l.reload(), "no-change reload must not error") + ptrAfter := l.Get() + assert.Same(t, ptrBefore, ptrAfter, "table pointer unchanged when mtime unchanged") +} + +func TestLoader_ReloadDetectsChange(t *testing.T) { + base := t.TempDir() + target := filepath.Join(base, "pricing.yaml") + copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target) + + l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil) + require.NoError(t, err) + + updated := []byte("openai:\n gpt-4o-mini:\n input_per_1k: 1.00\n output_per_1k: 2.00\n") + require.NoError(t, os.WriteFile(target, updated, 0o600)) + future := time.Now().Add(time.Hour) + require.NoError(t, os.Chtimes(target, future, future)) + + require.NoError(t, l.reload(), "reload must succeed on valid new content") + + cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0) + require.True(t, ok, "updated model still present") + assert.InDelta(t, 3.0, cost, 0.0001, "new prices are applied: 1 + 2 per 1k") +} + +// TestLoader_ReloadGoroutinePicksUpChanges proves the background goroutine +// started via Reload actually swaps the pricing table when the file changes +// on disk. Without that goroutine running, pricing edits would never reach +// requests until a proxy restart. +func TestLoader_ReloadGoroutinePicksUpChanges(t *testing.T) { + base := t.TempDir() + target := filepath.Join(base, "pricing.yaml") + copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target) + + l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil) + require.NoError(t, err) + l.SetReloadInterval(20 * time.Millisecond) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan struct{}) + go func() { + l.Reload(ctx) + close(done) + }() + + // Before any rewrite, the loader holds the fixture's prices. + costBefore, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0) + require.True(t, ok, "fixture model must resolve initially") + assert.InDelta(t, 0.00075, costBefore, 1e-9, "fixture prices apply before rewrite") + + updated := []byte("openai:\n gpt-4o-mini:\n input_per_1k: 1.00\n output_per_1k: 2.00\n") + require.NoError(t, os.WriteFile(target, updated, 0o600)) + future := time.Now().Add(time.Hour) + require.NoError(t, os.Chtimes(target, future, future)) + + deadline := time.Now().Add(2 * time.Second) + for { + cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0) + if ok && cost > 2.5 { + break + } + if time.Now().After(deadline) { + t.Fatalf("background reloader did not pick up rewrite within deadline") + } + time.Sleep(10 * time.Millisecond) + } + + cancel() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("Reload loop did not exit after cancel") + } +} + +func TestLoader_ReloadBackgroundLoopCancellation(t *testing.T) { + base := t.TempDir() + copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml")) + l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + l.Reload(ctx) + close(done) + }() + cancel() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("Reload loop did not exit on context cancel") + } +} + +func TestNewLoader_PathValidation(t *testing.T) { + base := t.TempDir() + copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml")) + + cases := []struct { + name string + relPath string + }{ + {"traversal", "../../etc/passwd"}, + {"absolute", "/etc/passwd"}, + {"slash in basename", "sub/pricing.yaml"}, + {"control chars", "pricing\x00.yaml"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + _, err := NewLoader(base, tc.relPath, "llm_observability", nil) + require.Error(t, err, "NewLoader must reject %q", tc.relPath) + }) + } + + // Empty relPath is no longer a validation error: the loader treats it + // as "no override file, defaults only" so cost metadata is still + // emitted for the embedded models out of the box. + t.Run("empty falls back to defaults", func(t *testing.T) { + l, err := NewLoader(base, "", "llm_observability", nil) + require.NoError(t, err, "empty relPath should yield a defaults-only loader") + require.NotNil(t, l, "loader must be returned") + require.False(t, l.WatchesFile(), "no file watching when no override is given") + _, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0) + assert.True(t, ok, "embedded defaults should still resolve gpt-4o-mini") + }) +} + +// TestNewLoader_PathValidation_Extended covers the remaining attack shapes +// called out in C2: dot references, embedded traversal segments, and a +// newline in the basename. The basename regex must reject each one even +// though filepath.Clean would otherwise collapse them. +func TestNewLoader_PathValidation_Extended(t *testing.T) { + base := t.TempDir() + copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml")) + + cases := []struct { + name string + relPath string + }{ + {"dot", "."}, + {"dotdot", ".."}, + {"relative traversal", "../pricing.yaml"}, + {"embedded slash", "pri/cing.yaml"}, + {"newline", "pricing\n.yaml"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + _, err := NewLoader(base, tc.relPath, "llm_observability", nil) + require.Error(t, err, "NewLoader must reject %q", tc.relPath) + }) + } +} + +// TestNewLoader_ValidBasenameLoads proves the allowlist is exclusive: a +// basename containing only safe characters under baseDir loads. Without this +// a regression that over-tightened the regex would silently break valid +// deployments. +func TestNewLoader_ValidBasenameLoads(t *testing.T) { + base := t.TempDir() + copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing-v2_prod.yaml")) + + l, err := NewLoader(base, "pricing-v2_prod.yaml", "llm_observability", nil) + require.NoError(t, err, "basename with _, -, . must load") + require.NotNil(t, l.Get(), "table populated") +} + +// TestNewLoader_SymlinkOutsideBaseDirRejected constructs a symlink under +// baseDir that points to a file outside it. O_NOFOLLOW must refuse to open +// the symlink even though the symlink path itself is a valid basename under +// baseDir. +func TestNewLoader_SymlinkOutsideBaseDirRejected(t *testing.T) { + outside := t.TempDir() + target := filepath.Join(outside, "evil.yaml") + copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target) + + base := t.TempDir() + link := filepath.Join(base, "pricing.yaml") + require.NoError(t, os.Symlink(target, link), "symlink setup") + + _, err := NewLoader(base, "pricing.yaml", "llm_observability", nil) + require.Error(t, err, "O_NOFOLLOW must reject symlink even when it points outside baseDir") +} + +func TestNewLoader_SymlinkRejected(t *testing.T) { + base := t.TempDir() + concrete := filepath.Join(base, "real.yaml") + copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), concrete) + + link := filepath.Join(base, "pricing.yaml") + require.NoError(t, os.Symlink(concrete, link), "symlink setup") + + _, err := NewLoader(base, "pricing.yaml", "llm_observability", nil) + require.Error(t, err, "O_NOFOLLOW must reject symlinked targets") +} + +func TestTableCost_NilSafe(t *testing.T) { + var t1 *Table + cost, ok := t1.Cost("x", "y", 1, 1, 0, 0) + assert.False(t, ok, "nil table reports unknown") + assert.Zero(t, cost, "nil table returns zero cost") + assert.False(t, t1.Has("x", "y"), "nil table has nothing") +} + +func TestLoaderGet_NilSafe(t *testing.T) { + var l *Loader + assert.Nil(t, l.Get(), "nil loader returns nil table") +} + +// TestNewLoader_RejectsOversizedFile_FixesM4 proves the loader bounds reads +// at maxPricingBytes so a hostile file cannot exhaust process memory. +func TestNewLoader_RejectsOversizedFile_FixesM4(t *testing.T) { + base := t.TempDir() + target := filepath.Join(base, "pricing.yaml") + + // Build a YAML payload larger than the cap. We pad with valid YAML + // comments so a partial read would still fail the size check rather + // than the parser. + header := "openai:\n" + bigComment := make([]byte, maxPricingBytes+1024) + for i := range bigComment { + bigComment[i] = ' ' + } + bigComment[0] = '#' + bigComment[len(bigComment)-1] = '\n' + payload := append([]byte(header), bigComment...) + require.NoError(t, os.WriteFile(target, payload, 0o600)) + + _, err := NewLoader(base, "pricing.yaml", "llm_observability", nil) + require.Error(t, err, "oversized pricing file must be rejected") + assert.Contains(t, err.Error(), "exceeds", "rejection must reference the byte cap") +} diff --git a/proxy/internal/llm/pricing/pricing_unix.go b/proxy/internal/llm/pricing/pricing_unix.go new file mode 100644 index 000000000..4f3ea33a2 --- /dev/null +++ b/proxy/internal/llm/pricing/pricing_unix.go @@ -0,0 +1,68 @@ +//go:build unix + +package pricing + +import ( + "fmt" + "io" + "os" + "syscall" + "time" + + log "github.com/sirupsen/logrus" +) + +// maxPricingBytes caps the size of the pricing YAML on read so a hostile or +// runaway file cannot exhaust process memory during reload. 1 MiB is several +// orders of magnitude larger than any reasonable pricing table. +const maxPricingBytes int64 = 1 << 20 + +// loadPricing opens the file with O_NOFOLLOW, fstats the open descriptor, +// and parses from that same descriptor. Never re-opens by path so a +// mid-read rename or symlink swap cannot substitute content. Bytes are +// capped at maxPricingBytes so the loader cannot be coerced into reading an +// unbounded file. +func loadPricing(path string) (*Table, time.Time, error) { + f, err := os.OpenFile(path, os.O_RDONLY|syscall.O_NOFOLLOW, 0) + if err != nil { + return nil, time.Time{}, fmt.Errorf("open %s: %w", path, err) + } + defer func() { + if cerr := f.Close(); cerr != nil { + log.Debugf("close pricing file %s: %v", path, cerr) + } + }() + + info, err := f.Stat() + if err != nil { + return nil, time.Time{}, fmt.Errorf("fstat %s: %w", path, err) + } + if !info.Mode().IsRegular() { + return nil, time.Time{}, fmt.Errorf("pricing file %s is not a regular file", path) + } + + data, err := io.ReadAll(io.LimitReader(f, maxPricingBytes+1)) + if err != nil { + return nil, time.Time{}, fmt.Errorf("read %s: %w", path, err) + } + if int64(len(data)) > maxPricingBytes { + return nil, time.Time{}, fmt.Errorf("pricing file %s exceeds %d bytes", path, maxPricingBytes) + } + + table, err := parsePricingBytes(data) + if err != nil { + return nil, time.Time{}, err + } + return table, info.ModTime(), nil +} + +// statMtime returns the mtime of the file at path. It uses lstat semantics +// via os.Lstat so a symlink swap is detected even though O_NOFOLLOW will +// later reject the open. +func statMtime(path string) (time.Time, error) { + info, err := os.Lstat(path) + if err != nil { + return time.Time{}, fmt.Errorf("lstat %s: %w", path, err) + } + return info.ModTime(), nil +} diff --git a/proxy/internal/llm/sse.go b/proxy/internal/llm/sse.go new file mode 100644 index 000000000..3d33ab577 --- /dev/null +++ b/proxy/internal/llm/sse.go @@ -0,0 +1,117 @@ +package llm + +import ( + "bufio" + "errors" + "fmt" + "io" + "strings" +) + +// Event represents a single server-sent event. Type is the dispatch name +// carried on an "event:" line (empty when the stream uses only "data:" +// lines). Data is the concatenation of every "data:" line that made up the +// event, joined by a single newline. +type Event struct { + Type string + Data string +} + +// Scanner reads SSE events from an underlying byte stream. Events are +// delimited by a blank line ("\n\n"). CRLF line endings are normalized to LF +// transparently so fixtures captured from live servers can be replayed. +// +// Scanner is not safe for concurrent use. +type Scanner struct { + r *bufio.Reader + maxLine int +} + +// NewScanner wraps the given reader. The default underlying buffer size is +// large enough for typical provider events (~64 KiB); callers needing +// larger events can wrap the reader in their own bufio.Reader beforehand. +func NewScanner(r io.Reader) *Scanner { + return &Scanner{ + r: bufio.NewReaderSize(r, 64*1024), + maxLine: 1 << 20, + } +} + +// Next returns the next event. It returns io.EOF after the final event has +// been consumed. A trailing event that is not terminated by a blank line is +// still returned before io.EOF so that servers which close the connection +// without a trailing newline are handled correctly. +func (s *Scanner) Next() (Event, error) { + var ( + event Event + dataBuf strings.Builder + hasData bool + hasAny bool + ) + + for { + line, err := s.readLine() + if err != nil { + if errors.Is(err, io.EOF) && hasAny { + event.Data = dataBuf.String() + return event, nil + } + return Event{}, err + } + + if line == "" { + if !hasAny { + continue + } + event.Data = dataBuf.String() + return event, nil + } + + hasAny = true + if strings.HasPrefix(line, ":") { + continue + } + + field, value := splitField(line) + switch field { + case "event": + event.Type = value + case "data": + if hasData { + dataBuf.WriteByte('\n') + } + dataBuf.WriteString(value) + hasData = true + } + } +} + +func (s *Scanner) readLine() (string, error) { + line, err := s.r.ReadString('\n') + if err != nil { + if errors.Is(err, io.EOF) && line != "" { + return trimEOL(line), nil + } + return "", err + } + if len(line) > s.maxLine { + return "", fmt.Errorf("sse line exceeds %d bytes", s.maxLine) + } + return trimEOL(line), nil +} + +func trimEOL(line string) string { + line = strings.TrimRight(line, "\n") + line = strings.TrimRight(line, "\r") + return line +} + +func splitField(line string) (string, string) { + idx := strings.IndexByte(line, ':') + if idx < 0 { + return line, "" + } + field := line[:idx] + value := strings.TrimPrefix(line[idx+1:], " ") + return field, value +} diff --git a/proxy/internal/llm/sse_test.go b/proxy/internal/llm/sse_test.go new file mode 100644 index 000000000..96cecc111 --- /dev/null +++ b/proxy/internal/llm/sse_test.go @@ -0,0 +1,175 @@ +package llm + +import ( + "errors" + "io" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func collectEvents(t *testing.T, r io.Reader) []Event { + t.Helper() + s := NewScanner(r) + var out []Event + for { + ev, err := s.Next() + if errors.Is(err, io.EOF) { + return out + } + require.NoError(t, err, "unexpected error scanning SSE") + out = append(out, ev) + } +} + +func TestSSEScanner_OpenAIFixture(t *testing.T) { + f, err := os.Open(filepath.Join("fixtures", "openai_stream.txt")) + require.NoError(t, err, "fixture must be openable") + defer f.Close() + + events := collectEvents(t, f) + require.Len(t, events, 4, "expected 4 data frames (3 chunks + [DONE])") + + for _, ev := range events { + assert.Empty(t, ev.Type, "OpenAI stream uses data-only frames") + } + assert.Contains(t, events[2].Data, `"usage"`, "third chunk carries usage block") + assert.Equal(t, "[DONE]", events[3].Data, "final frame is the OpenAI DONE sentinel") +} + +func TestSSEScanner_AnthropicFixture(t *testing.T) { + f, err := os.Open(filepath.Join("fixtures", "anthropic_stream.txt")) + require.NoError(t, err, "fixture must be openable") + defer f.Close() + + events := collectEvents(t, f) + require.Len(t, events, 7, "expected 7 Anthropic events") + + types := make([]string, 0, len(events)) + for _, ev := range events { + types = append(types, ev.Type) + } + assert.Equal(t, []string{ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + }, types, "Anthropic event ordering matches fixture") + + var deltaUsage Event + for _, ev := range events { + if ev.Type == "message_delta" { + deltaUsage = ev + break + } + } + assert.Contains(t, deltaUsage.Data, `"output_tokens":45`, "message_delta carries partial usage") +} + +func TestSSEScanner_MultilineData(t *testing.T) { + raw := "event: ping\ndata: line1\ndata: line2\ndata: line3\n\n" + events := collectEvents(t, strings.NewReader(raw)) + + require.Len(t, events, 1, "one logical event from three data lines") + assert.Equal(t, "ping", events[0].Type, "event name honored") + assert.Equal(t, "line1\nline2\nline3", events[0].Data, "data lines joined with newline") +} + +func TestSSEScanner_CRLF(t *testing.T) { + raw := "event: foo\r\ndata: bar\r\n\r\ndata: baz\r\n\r\n" + events := collectEvents(t, strings.NewReader(raw)) + + require.Len(t, events, 2, "CRLF-delimited events recognized") + assert.Equal(t, "foo", events[0].Type, "first event type preserved") + assert.Equal(t, "bar", events[0].Data, "first event data preserved") + assert.Empty(t, events[1].Type, "second event has no event name") + assert.Equal(t, "baz", events[1].Data, "second event data preserved") +} + +func TestSSEScanner_EmptyInput(t *testing.T) { + s := NewScanner(strings.NewReader("")) + _, err := s.Next() + require.ErrorIs(t, err, io.EOF, "empty input yields immediate EOF") +} + +func TestSSEScanner_CommentIgnored(t *testing.T) { + raw := ": this is a comment\ndata: hi\n\n" + events := collectEvents(t, strings.NewReader(raw)) + require.Len(t, events, 1, "comment line does not emit an event") + assert.Equal(t, "hi", events[0].Data, "data line honoured after comment") +} + +func TestSSEScanner_TrailingWithoutBlankLine(t *testing.T) { + raw := "event: foo\ndata: bar\n" + events := collectEvents(t, strings.NewReader(raw)) + require.Len(t, events, 1, "trailing event without blank line still emitted") + assert.Equal(t, "foo", events[0].Type) + assert.Equal(t, "bar", events[0].Data) +} + +// TestSSEScanner_ManyConsecutiveEmptyLines feeds a stream that is nothing +// but empty lines. The scanner must terminate without panic — empty lines +// alone do not constitute an event and must yield io.EOF. +func TestSSEScanner_ManyConsecutiveEmptyLines(t *testing.T) { + raw := strings.Repeat("\n", 100) + s := NewScanner(strings.NewReader(raw)) + _, err := s.Next() + require.ErrorIs(t, err, io.EOF, "100 empty lines must terminate as EOF without panic") +} + +// TestSSEScanner_InterleavedCRLFAndLF mixes \r\n and \n terminators within +// the same event. The scanner normalizes both and must still recover a +// coherent event. +func TestSSEScanner_InterleavedCRLFAndLF(t *testing.T) { + raw := "event: mix\r\ndata: first\ndata: second\r\n\n" + events := collectEvents(t, strings.NewReader(raw)) + require.Len(t, events, 1, "mixed line endings must still produce one event") + assert.Equal(t, "mix", events[0].Type) + assert.Equal(t, "first\nsecond", events[0].Data, "both data lines joined") +} + +// TestSSEScanner_LongSingleDataLine constructs a single data line that +// exceeds the default bufio buffer (64 KiB) but stays under the scanner +// maxLine. The scanner must round-trip the value intact without panicking +// or truncating silently. +func TestSSEScanner_LongSingleDataLine(t *testing.T) { + big := strings.Repeat("x", 80<<10) + raw := "data: " + big + "\n\n" + events := collectEvents(t, strings.NewReader(raw)) + require.Len(t, events, 1, "long single-line event must be emitted") + assert.Equal(t, big, events[0].Data, "long data preserved") +} + +// TestSSEScanner_BinaryGarbageInData validates that non-printable bytes +// inside a data line do not crash the parser. The scanner should either +// round-trip them or return a well-formed error — never panic. +func TestSSEScanner_BinaryGarbageInData(t *testing.T) { + raw := "data: \x00\x01\x02\xff\xfe\n\n" + defer func() { + if r := recover(); r != nil { + t.Fatalf("scanner panicked on binary garbage: %v", r) + } + }() + s := NewScanner(strings.NewReader(raw)) + ev, err := s.Next() + require.NoError(t, err, "binary bytes in data should not surface as error") + assert.Equal(t, "\x00\x01\x02\xff\xfe", ev.Data, "binary payload round-trips") +} + +// TestSSEScanner_UnknownFieldsIgnored stresses the field parser by sending +// unrecognized field names ("id:", "retry:", "custom:"). They must be +// silently ignored per the SSE spec; the scanner must not panic or emit +// spurious events. +func TestSSEScanner_UnknownFieldsIgnored(t *testing.T) { + raw := "id: 1\nretry: 5000\ncustom: value\ndata: payload\n\n" + events := collectEvents(t, strings.NewReader(raw)) + require.Len(t, events, 1, "unknown fields must not spawn extra events") + assert.Equal(t, "payload", events[0].Data, "data field survives amid unknown fields") +} diff --git a/proxy/internal/middleware/builtin/agentnetwork_chain_integration_test.go b/proxy/internal/middleware/builtin/agentnetwork_chain_integration_test.go new file mode 100644 index 000000000..643990876 --- /dev/null +++ b/proxy/internal/middleware/builtin/agentnetwork_chain_integration_test.go @@ -0,0 +1,318 @@ +package builtin_test + +import ( + "context" + "net" + "runtime" + "strconv" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/test/bufconn" + + mgmtgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc" + "github.com/netbirdio/netbird/management/server/agentnetwork" + agentNetworkTypes "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/store" + nbtypes "github.com/netbirdio/netbird/management/server/types" + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_check" + "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_record" + "github.com/netbirdio/netbird/shared/management/proto" +) + +// chainIntegrationFixture wires the BOTH new agent-network +// middlewares against a live in-process management stack: real +// sqlite store + real Manager + real gRPC server. The proxy chain +// framework itself isn't constructed (its dispatcher / accumulator / +// metadata gate are tested separately); we exercise the middleware +// pair as the proxy runtime would, by invoking each with a crafted +// Input and asserting the wire path between them. +// +// This is the regression cover for item 16 in the design review: +// real LLM request → cost stamped → consumption row in the table. +type chainIntegrationFixture struct { + store store.Store + manager agentnetwork.Manager + gatecase *llm_limit_check.Middleware + recorder *llm_limit_record.Middleware +} + +func newChainIntegration(t *testing.T) *chainIntegrationFixture { + t.Helper() + if runtime.GOOS == "windows" { + t.Skip("sqlite store not properly supported on Windows yet") + } + t.Setenv("NETBIRD_STORE_ENGINE", string(nbtypes.SqliteStoreEngine)) + + st, cleanUp, err := store.NewTestStoreFromSQL(context.Background(), "", t.TempDir()) + require.NoError(t, err) + t.Cleanup(cleanUp) + + manager := agentnetwork.NewManager(st, nil, nil, nil) + + server := &mgmtgrpc.ProxyServiceServer{} + server.SetAgentNetworkLimitsService(manager) + + const bufSize = 1024 * 1024 + lis := bufconn.Listen(bufSize) + srv := grpc.NewServer() + proto.RegisterProxyServiceServer(srv, server) + go func() { _ = srv.Serve(lis) }() + t.Cleanup(srv.Stop) + + conn, err := grpc.NewClient("passthrough:///bufnet", + grpc.WithContextDialer(func(_ context.Context, _ string) (net.Conn, error) { return lis.Dial() }), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + + mgmtClient := proto.NewProxyServiceClient(conn) + return &chainIntegrationFixture{ + store: st, + manager: manager, + gatecase: llm_limit_check.New(mgmtClient, nil), + recorder: llm_limit_record.New(mgmtClient, nil), + } +} + +// chainInput builds a middleware Input that mirrors what the proxy +// framework would synthesise for a tunnel-peer LLM request. The +// gate consumes the resolved provider id from upstream metadata +// (set by llm_router); the recorder consumes the attribution +// metadata stamped by the gate plus tokens / cost from +// llm_response_parser + cost_meter. +func chainInput(account, user, group, providerID string, requestMeta []middleware.KV) *middleware.Input { + _ = providerID // packed into requestMeta by the caller as KeyLLMResolvedProviderID + return &middleware.Input{ + AccountID: account, + UserID: user, + UserGroups: []string{group}, + Metadata: requestMeta, + } +} + +// chainCapPolicy builds a tight token-cap policy fixture for the +// chain integration tests. Inlined here (rather than imported) because +// the equivalent helper in the management gRPC package is unexported +// and this is a different package boundary. +func chainCapPolicy(id, account string, sourceGroups []string, providerID string, tokenCap, windowSec int64) *agentNetworkTypes.Policy { + return &agentNetworkTypes.Policy{ + ID: id, + AccountID: account, + Enabled: true, + Name: id, + SourceGroups: sourceGroups, + DestinationProviderIDs: []string{providerID}, + Limits: agentNetworkTypes.PolicyLimits{ + TokenLimit: agentNetworkTypes.PolicyTokenLimit{ + Enabled: true, + GroupCap: tokenCap, + WindowSeconds: windowSec, + }, + }, + CreatedAt: time.Now().UTC(), + UpdatedAt: time.Now().UTC(), + } +} + +// TestChain_AllowPath_StampsAttributionAndRecordsCounter walks the +// full happy path: gate calls CheckLLMPolicyLimits → stamps +// attribution metadata → recorder reads metadata + tokens / cost → +// calls RecordLLMUsage → counters land in sqlite. Asserting on the +// store at the end proves every leg of the wire works together, +// not just each leg in isolation (which the unit tests already cover). +func TestChain_AllowPath_StampsAttributionAndRecordsCounter(t *testing.T) { + f := newChainIntegration(t) + + const account = "acc-1" + const user = "user-bob" + const group = "grp-engineers" + const provider = "prov-1" + + // Seed a policy with token + budget caps; both halves carry + // real ceilings so the request stays within headroom. + require.NoError(t, f.store.SaveAgentNetworkPolicy(context.Background(), + chainCapPolicy("pol-1", account, []string{group}, provider, 10_000, 86_400))) + + // ── Stage 1 — gate: pre-flight check ────────────────────── + gateIn := chainInput(account, user, group, provider, []middleware.KV{ + {Key: middleware.KeyLLMResolvedProviderID, Value: provider}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + }) + gateOut, err := f.gatecase.Invoke(context.Background(), gateIn) + require.NoError(t, err) + require.Equal(t, middleware.DecisionAllow, gateOut.Decision, "fresh policy must allow") + + // Verify attribution metadata was stamped — the recorder + // depends on these keys. + metaMap := map[string]string{} + for _, kv := range gateOut.Metadata { + metaMap[kv.Key] = kv.Value + } + assert.Equal(t, "pol-1", metaMap[middleware.KeyLLMSelectedPolicyID]) + assert.Equal(t, group, metaMap[middleware.KeyLLMAttributionGroupID]) + assert.Equal(t, "86400", metaMap[middleware.KeyLLMAttributionWindowS]) + + // ── Stage 2 — recorder: post-flight write ───────────────── + // Build the response-leg Input the framework would synthesise + // for the recorder: gate's emitted attribution metadata + the + // tokens / cost stamped by llm_response_parser + cost_meter. + const tokensIn = int64(123) + const tokensOut = int64(45) + const costUSD = 0.0042 + recordIn := chainInput(account, user, group, provider, append([]middleware.KV{}, + gateOut.Metadata...)) + recordIn.Metadata = append(recordIn.Metadata, + middleware.KV{Key: middleware.KeyLLMInputTokens, Value: strconv.FormatInt(tokensIn, 10)}, + middleware.KV{Key: middleware.KeyLLMOutputTokens, Value: strconv.FormatInt(tokensOut, 10)}, + middleware.KV{Key: middleware.KeyCostUSDTotal, Value: strconv.FormatFloat(costUSD, 'f', 6, 64)}, + ) + recordOut, err := f.recorder.Invoke(context.Background(), recordIn) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, recordOut.Decision, "recorder always allows; its only side effect is the counter write") + + // ── Stage 3 — assert state in sqlite ────────────────────── + windowStart := agentNetworkTypes.WindowStart(time.Now(), 86_400) + userRow, err := f.store.GetAgentNetworkConsumption( + context.Background(), store.LockingStrengthNone, account, + agentNetworkTypes.DimensionUser, user, int64(86_400), windowStart, + ) + require.NoError(t, err) + assert.Equal(t, tokensIn, userRow.TokensInput, "user counter must hold the input tokens the recorder posted") + assert.Equal(t, tokensOut, userRow.TokensOutput) + assert.InDelta(t, costUSD, userRow.CostUSD, 1e-6) + + groupRow, err := f.store.GetAgentNetworkConsumption( + context.Background(), store.LockingStrengthNone, account, + agentNetworkTypes.DimensionGroup, group, int64(86_400), windowStart, + ) + require.NoError(t, err) + assert.Equal(t, tokensIn, groupRow.TokensInput, "group counter mirrors the user counter — single Record posts both dims") +} + +// TestChain_DenyPath_GateRejectsAndNoConsumptionWritten covers the +// negative side: when the gate denies, the recorder is never +// invoked (the proxy framework short-circuits on Decision=Deny). +// We assert no consumption row materialises after the gate-deny +// path, even though the test technically calls the recorder +// afterwards — the recorder must skip on missing attribution +// metadata so the framework's short-circuit isn't load-bearing for +// data integrity. +func TestChain_DenyPath_GateRejectsAndNoConsumptionWritten(t *testing.T) { + f := newChainIntegration(t) + + const account = "acc-1" + const user = "user-bob" + const group = "grp-tight" + const provider = "prov-1" + + policy := chainCapPolicy("pol-tight", account, []string{group}, provider, 100, 86_400) + require.NoError(t, f.store.SaveAgentNetworkPolicy(context.Background(), policy)) + + // Pre-burn the counter to the cap so the gate denies. + require.NoError(t, f.store.IncrementAgentNetworkConsumption( + context.Background(), account, + agentNetworkTypes.DimensionGroup, group, int64(86_400), + agentNetworkTypes.WindowStart(time.Now(), 86_400), + 100, 0, 0, + )) + + gateOut, err := f.gatecase.Invoke(context.Background(), chainInput(account, user, group, provider, + []middleware.KV{ + {Key: middleware.KeyLLMResolvedProviderID, Value: provider}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + }, + )) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionDeny, gateOut.Decision, "policy at-cap must deny on the gate") + require.NotNil(t, gateOut.DenyReason) + assert.Equal(t, "llm_policy.token_cap_exceeded", gateOut.DenyReason.Code) + + // On deny, the gate emits no attribution metadata. If the + // proxy framework still invokes the recorder (defense in + // depth), the recorder's "no attribution window = skip" guard + // prevents a phantom counter increment. + recordOut, err := f.recorder.Invoke(context.Background(), chainInput(account, user, group, provider, + gateOut.Metadata, // no llm.attribution_window_seconds stamped + )) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, recordOut.Decision) + + // The pre-burned 100 tokens are the only counter movement — + // the recorder must NOT have added a fresh row for the user + // dimension on this denied request. + windowStart := agentNetworkTypes.WindowStart(time.Now(), 86_400) + userRow, err := f.store.GetAgentNetworkConsumption( + context.Background(), store.LockingStrengthNone, account, + agentNetworkTypes.DimensionUser, user, int64(86_400), windowStart, + ) + require.NoError(t, err) + assert.Zero(t, userRow.TokensInput, "user dimension must not gain tokens from a denied request — recorder skip is the safety net") +} + +// TestChain_CapExhaustTransition exercises the allow→deny boundary +// the operator cares most about: a request just under cap allows +// AND records, the next request post-record at-cap denies. This is +// the same lifecycle 50-grpc-allow-record-deny.sh runs in bash, but +// against the actual middleware pair rather than the smoke binary +// driving the gRPC RPCs directly. +func TestChain_CapExhaustTransition(t *testing.T) { + f := newChainIntegration(t) + + const account = "acc-1" + const user = "user-alice" + const group = "grp-cap-edge" + const provider = "prov-1" + const tightCap = int64(100) + + require.NoError(t, f.store.SaveAgentNetworkPolicy(context.Background(), + chainCapPolicy("pol-edge", account, []string{group}, provider, tightCap, 86_400))) + + // Pre-burn 99 tokens so we're at the very edge. + require.NoError(t, f.store.IncrementAgentNetworkConsumption( + context.Background(), account, + agentNetworkTypes.DimensionGroup, group, int64(86_400), + agentNetworkTypes.WindowStart(time.Now(), 86_400), + 99, 0, 0, + )) + + // Gate at 99/100 — must allow (one token of headroom). + gateOut, err := f.gatecase.Invoke(context.Background(), chainInput(account, user, group, provider, + []middleware.KV{ + {Key: middleware.KeyLLMResolvedProviderID, Value: provider}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + }, + )) + require.NoError(t, err) + require.Equal(t, middleware.DecisionAllow, gateOut.Decision, "99/100 must allow — one token of headroom") + + // Record one more input token — pushes us to 100/100. + recordIn := chainInput(account, user, group, provider, append([]middleware.KV{}, + gateOut.Metadata...)) + recordIn.Metadata = append(recordIn.Metadata, + middleware.KV{Key: middleware.KeyLLMInputTokens, Value: "1"}, + middleware.KV{Key: middleware.KeyLLMOutputTokens, Value: "0"}, + middleware.KV{Key: middleware.KeyCostUSDTotal, Value: "0.000001"}, + ) + _, err = f.recorder.Invoke(context.Background(), recordIn) + require.NoError(t, err) + + // Next gate call must deny — counter is exactly at cap. + gateOut2, err := f.gatecase.Invoke(context.Background(), chainInput(account, user, group, provider, + []middleware.KV{ + {Key: middleware.KeyLLMResolvedProviderID, Value: provider}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + }, + )) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionDeny, gateOut2.Decision, + "once recorder pushed the group counter to 100/100, the next gate call must deny — allow→deny transition is the operator-visible product semantic") + require.NotNil(t, gateOut2.DenyReason) + assert.Equal(t, "llm_policy.token_cap_exceeded", gateOut2.DenyReason.Code) +} diff --git a/proxy/internal/middleware/builtin/all_test.go b/proxy/internal/middleware/builtin/all_test.go new file mode 100644 index 000000000..28576e248 --- /dev/null +++ b/proxy/internal/middleware/builtin/all_test.go @@ -0,0 +1,40 @@ +package builtin_test + +import ( + "sort" + "testing" + + "github.com/stretchr/testify/assert" + + mwbuiltin "github.com/netbirdio/netbird/proxy/internal/middleware/builtin" + + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/cost_meter" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_guardrail" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_identity_inject" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_check" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_record" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_request_parser" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_response_parser" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_router" +) + +// TestDefaultRegistry_BuiltinIDs locks the set of middleware IDs that +// the default builtin registry exposes once every sub-package's init() +// has run. The list is the source of truth wired by the synthesiser +// in management; adding a new built-in middleware should consciously +// extend this list. +func TestDefaultRegistry_BuiltinIDs(t *testing.T) { + got := mwbuiltin.DefaultRegistry().IDs() + sort.Strings(got) + want := []string{ + "cost_meter", + "llm_guardrail", + "llm_identity_inject", + "llm_limit_check", + "llm_limit_record", + "llm_request_parser", + "llm_response_parser", + "llm_router", + } + assert.Equal(t, want, got, "default registry must expose every built-in middleware after anonymous imports") +} diff --git a/proxy/internal/middleware/builtin/builtin.go b/proxy/internal/middleware/builtin/builtin.go new file mode 100644 index 000000000..9ea4cf89d --- /dev/null +++ b/proxy/internal/middleware/builtin/builtin.go @@ -0,0 +1,93 @@ +// Package builtin holds the package-level middleware registry that +// concrete middleware packages register themselves into via init(). +// Server boot anonymous-imports each middleware sub-package; the +// resolver attached to the middleware Manager pulls factories out of +// this registry. +package builtin + +import ( + "context" + "sync" + + log "github.com/sirupsen/logrus" + "go.opentelemetry.io/otel/metric" + "google.golang.org/grpc" + + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/shared/management/proto" +) + +// MgmtClient is the narrow slice of proto.ProxyServiceClient that +// builtin middlewares may use during request / response handling. +// Only the agent-network limit pair (llm_limit_check + llm_limit_record) +// uses this today; declaring the surface here keeps the dependency +// explicit at boot time. +// +// proto.ProxyServiceClient already satisfies this interface so server +// boot just forwards its existing client. +type MgmtClient interface { + CheckLLMPolicyLimits(ctx context.Context, in *proto.CheckLLMPolicyLimitsRequest, opts ...grpc.CallOption) (*proto.CheckLLMPolicyLimitsResponse, error) + RecordLLMUsage(ctx context.Context, in *proto.RecordLLMUsageRequest, opts ...grpc.CallOption) (*proto.RecordLLMUsageResponse, error) +} + +// defaultRegistry is the package-level registry that concrete builtin +// middlewares register themselves into via init(). +var defaultRegistry = middleware.NewRegistry() + +// FactoryContext is the per-process bag that concrete factories may +// consult during construction. It carries the proxy-lifetime context, +// the data directory used for static config files (pricing tables, +// allowlists), the OTel meter, and the proxy logger. +// +// Configure must be called once at boot before any chain build calls +// Resolve. Calling it twice overwrites the prior value; tests may rely +// on this to reset state. +type FactoryContext struct { + Context context.Context + DataDir string + Meter metric.Meter + Logger *log.Logger + MgmtClient MgmtClient +} + +var ( + ctxStore FactoryContext + ctxMu sync.RWMutex +) + +// Configure stores the per-process FactoryContext. Concrete factories +// reach for it via Context(). mgmt may be nil on tests / standalone +// builds with no management server; consumers must guard. +func Configure(ctx context.Context, dataDir string, meter metric.Meter, logger *log.Logger, mgmt MgmtClient) { + ctxMu.Lock() + defer ctxMu.Unlock() + ctxStore = FactoryContext{ + Context: ctx, + DataDir: dataDir, + Meter: meter, + Logger: logger, + MgmtClient: mgmt, + } +} + +// Context returns the stored FactoryContext. Returns a zero value when +// Configure was never called; consumers must guard against nil +// Context/Meter/Logger if they care. +func Context() FactoryContext { + ctxMu.RLock() + defer ctxMu.RUnlock() + return ctxStore +} + +// Register adds a factory to the default registry. Called from init() +// blocks of concrete middleware packages. Panics on collision so +// duplicate IDs surface at startup. +func Register(f middleware.Factory) { + defaultRegistry.MustRegister(f) +} + +// DefaultRegistry returns the shared registry. The proxy server +// constructs the Resolver from it at boot. +func DefaultRegistry() *middleware.Registry { + return defaultRegistry +} diff --git a/proxy/internal/middleware/builtin/cost_meter/factory.go b/proxy/internal/middleware/builtin/cost_meter/factory.go new file mode 100644 index 000000000..b8a58d10e --- /dev/null +++ b/proxy/internal/middleware/builtin/cost_meter/factory.go @@ -0,0 +1,88 @@ +package cost_meter + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + + "github.com/netbirdio/netbird/proxy/internal/llm/pricing" + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/proxy/internal/middleware/builtin" +) + +// defaultPricingFilename is the basename probed inside the proxy data +// directory when no override is configured. +const defaultPricingFilename = "pricing.yaml" + +// Config is the on-wire configuration for the middleware. +type Config struct { + // PricingPath optionally overrides the basename of the pricing + // file probed inside the proxy data directory. When empty the + // loader falls back to "pricing.yaml". + PricingPath string `json:"pricing_path"` +} + +// Factory builds cost_meter 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. The +// pricing loader is built once per instance and reused across +// invocations. +func (Factory) New(rawConfig []byte) (middleware.Middleware, error) { + cfg, err := decodeConfig(rawConfig) + if err != nil { + return nil, err + } + + fctx := builtin.Context() + pricingPath := cfg.PricingPath + if pricingPath == "" { + pricingPath = defaultPricingFilename + } + + loader, err := pricing.NewLoader(fctx.DataDir, pricingPath, ID, nil) + if err != nil { + return nil, fmt.Errorf("init pricing loader: %w", err) + } + + cancel := startReloader(fctx.Context, loader) + + return newMiddleware(loader, cancel), nil +} + +// startReloader binds the loader's mtime-poll goroutine to a context +// derived from the proxy-lifetime context and returns its cancel func so +// the owning middleware can stop the goroutine on teardown. Returns nil +// when there's nothing to watch (nil context or defaults-only loader), in +// which case the middleware's Close is a no-op. +func startReloader(ctx context.Context, loader *pricing.Loader) context.CancelFunc { + if ctx == nil || !loader.WatchesFile() { + return nil + } + cctx, cancel := context.WithCancel(ctx) + go loader.Reload(cctx) + return cancel +} + +// decodeConfig accepts empty, null, and {} configs, returning a +// zero-value Config. Non-empty payloads must parse cleanly. +func decodeConfig(rawConfig []byte) (Config, error) { + var cfg Config + if len(bytes.TrimSpace(rawConfig)) == 0 { + return cfg, nil + } + if err := json.Unmarshal(rawConfig, &cfg); err != nil { + return cfg, fmt.Errorf("decode config: %w", err) + } + return cfg, nil +} + +func init() { + builtin.Register(Factory{}) +} diff --git a/proxy/internal/middleware/builtin/cost_meter/middleware.go b/proxy/internal/middleware/builtin/cost_meter/middleware.go new file mode 100644 index 000000000..4da620310 --- /dev/null +++ b/proxy/internal/middleware/builtin/cost_meter/middleware.go @@ -0,0 +1,193 @@ +// Package cost_meter implements the SlotOnResponse middleware that +// converts token-usage metadata emitted by llm_response_parser into a +// per-request USD cost estimate. The middleware uses the shared pricing +// loader so operator pricing overrides apply to the chain. +package cost_meter + +import ( + "context" + "fmt" + "strconv" + + "github.com/netbirdio/netbird/proxy/internal/llm/pricing" + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +// ID is the registry identifier for this middleware. +const ID = "cost_meter" + +// Version is the implementation version emitted via the spec merge. +const Version = "1.0.0" + +// Skip reasons emitted under KeyCostSkipped. The set is closed; the +// dashboard surfaces these verbatim. +const ( + skipMissingProvider = "missing_provider" + skipMissingModel = "missing_model" + skipMissingTokens = "missing_tokens" + //nolint:gosec // skip-reason label, not a credential + skipUnparseableTokens = "unparseable_tokens" + skipZeroTokens = "zero_tokens" + skipUnknownModel = "unknown_model" +) + +var metadataKeys = []string{ + middleware.KeyCostUSDTotal, + middleware.KeyCostSkipped, +} + +// Middleware computes a per-response cost estimate from the token +// counts emitted upstream by llm_response_parser. +type Middleware struct { + loader *pricing.Loader + // cancel stops this instance's pricing-reload goroutine. Non-nil only + // when the loader watches an override file; Close calls it so a chain + // rebuild doesn't leak a poll goroutine per retired instance. + cancel context.CancelFunc +} + +// newMiddleware constructs a Middleware bound to the given pricing loader. +// cancel may be nil (defaults-only loader with no reloader to stop). +func newMiddleware(loader *pricing.Loader, cancel context.CancelFunc) *Middleware { + return &Middleware{loader: loader, cancel: cancel} +} + +// 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 that the middleware runs after the upstream call. +func (m *Middleware) Slot() middleware.Slot { return middleware.SlotOnResponse } + +// AcceptedContentTypes is empty: cost_meter never inspects bodies. +func (m *Middleware) AcceptedContentTypes() []string { return []string{} } + +// MetadataKeys returns the closed allowlist of keys this middleware +// may emit. +func (m *Middleware) MetadataKeys() []string { + return append([]string(nil), metadataKeys...) +} + +// MutationsSupported reports that this middleware never mutates the +// response. +func (m *Middleware) MutationsSupported() bool { return false } + +// Close stops this instance's pricing-reload goroutine, if any. Called by +// the chain when a rebuild retires the instance, so the mtime-poll loop +// doesn't outlive the chain it belonged to. Safe to call on a nil receiver +// and on an instance with no reloader. +func (m *Middleware) Close() error { + if m != nil && m.cancel != nil { + m.cancel() + } + return nil +} + +// Invoke reads provider, model, and token metadata, looks up pricing, +// and emits either KeyCostUSDTotal or KeyCostSkipped. The decision is +// always DecisionAllow; cost metering never denies or mutates. +func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) { + out := &middleware.Output{Decision: middleware.DecisionAllow} + if in == nil { + return out, nil + } + + provider := lookupKV(in.Metadata, middleware.KeyLLMProvider) + if provider == "" { + out.Metadata = skip(skipMissingProvider) + return out, nil + } + + model := lookupKV(in.Metadata, middleware.KeyLLMModel) + if model == "" { + out.Metadata = skip(skipMissingModel) + return out, nil + } + + inRaw, hasIn := lookupKVOK(in.Metadata, middleware.KeyLLMInputTokens) + outRaw, hasOut := lookupKVOK(in.Metadata, middleware.KeyLLMOutputTokens) + if !hasIn || !hasOut { + out.Metadata = skip(skipMissingTokens) + return out, nil + } + + inTokens, err := strconv.ParseInt(inRaw, 10, 64) + if err != nil || inTokens < 0 { + // Unparseable or negative tokens are not a runtime error: the + // upstream llm_response_parser emitted a non-numeric / invalid + // value, so we surface that as cost.skipped and continue with + // Allow rather than pricing a negative count. + out.Metadata = skip(skipUnparseableTokens) + return out, nil //nolint:nilerr // structured skip; not a runtime error + } + outTokens, err := strconv.ParseInt(outRaw, 10, 64) + if err != nil || outTokens < 0 { + out.Metadata = skip(skipUnparseableTokens) + return out, nil //nolint:nilerr // structured skip; not a runtime error + } + + // Cache buckets are optional and silently zeroed on a missing / + // malformed value; they're a refinement on top of input cost, + // not a precondition. A buggy value falls back to 0, never aborts. + cachedTokens := parseOptionalInt64(in.Metadata, middleware.KeyLLMCachedInputTokens) + cacheCreationTokens := parseOptionalInt64(in.Metadata, middleware.KeyLLMCacheCreationTokens) + + if inTokens == 0 && outTokens == 0 && cachedTokens == 0 && cacheCreationTokens == 0 { + out.Metadata = skip(skipZeroTokens) + return out, nil + } + + table := m.loader.Get() + cost, ok := table.Cost(provider, model, inTokens, outTokens, cachedTokens, cacheCreationTokens) + if !ok { + out.Metadata = skip(skipUnknownModel) + return out, nil + } + + out.Metadata = []middleware.KV{ + {Key: middleware.KeyCostUSDTotal, Value: fmt.Sprintf("%.6f", cost)}, + } + return out, nil +} + +// skip returns a single-entry metadata slice carrying the given skip +// reason under KeyCostSkipped. +func skip(reason string) []middleware.KV { + return []middleware.KV{{Key: middleware.KeyCostSkipped, Value: reason}} +} + +// lookupKV returns the value associated with key, or the empty string +// when the key is absent. +func lookupKV(kvs []middleware.KV, key string) string { + v, _ := lookupKVOK(kvs, key) + return v +} + +// lookupKVOK returns the value associated with key plus a presence +// flag so callers can distinguish absent from empty. +func lookupKVOK(kvs []middleware.KV, key string) (string, bool) { + for _, kv := range kvs { + if kv.Key == key { + return kv.Value, true + } + } + return "", false +} + +// parseOptionalInt64 reads a metadata value and decodes it as int64. +// Absent or unparseable values yield 0 — the caller treats absence as +// "no cached tokens" rather than an error, since cache buckets are a +// refinement, not a precondition. +func parseOptionalInt64(kvs []middleware.KV, key string) int64 { + raw, ok := lookupKVOK(kvs, key) + if !ok { + return 0 + } + v, err := strconv.ParseInt(raw, 10, 64) + if err != nil || v < 0 { + return 0 + } + return v +} diff --git a/proxy/internal/middleware/builtin/cost_meter/middleware_test.go b/proxy/internal/middleware/builtin/cost_meter/middleware_test.go new file mode 100644 index 000000000..d1c161cab --- /dev/null +++ b/proxy/internal/middleware/builtin/cost_meter/middleware_test.go @@ -0,0 +1,459 @@ +package cost_meter + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/proxy/internal/middleware/builtin" +) + +const fixturePricing = `openai: + gpt-4o: + input_per_1k: 0.0025 + output_per_1k: 0.01 + gpt-4o-mini: + input_per_1k: 0.00015 + output_per_1k: 0.0006 +anthropic: + claude-sonnet-4-5: + input_per_1k: 0.003 + output_per_1k: 0.015 +` + +// configureBuiltin points the package-level FactoryContext at a tmp +// directory containing the test pricing fixture. Returns the path so +// callers can override files later if needed. +func configureBuiltin(t *testing.T) string { + t.Helper() + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "pricing.yaml"), []byte(fixturePricing), 0o600), "write pricing fixture") + builtin.Configure(context.Background(), dir, nil, nil, nil) + return dir +} + +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 buildMiddleware(t *testing.T, raw []byte) middleware.Middleware { + t.Helper() + mw, err := Factory{}.New(raw) + require.NoError(t, err, "factory must accept the supplied config") + return mw +} + +func TestMiddleware_StaticSurface(t *testing.T) { + configureBuiltin(t) + mw := buildMiddleware(t, nil) + + 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.SlotOnResponse, mw.Slot(), "must run in the response slot") + assert.Empty(t, mw.AcceptedContentTypes(), "cost_meter does not inspect bodies") + assert.False(t, mw.MutationsSupported(), "cost_meter never mutates") + assert.NoError(t, mw.Close(), "Close on stateless middleware is a no-op") + + keys := mw.MetadataKeys() + expected := []string{middleware.KeyCostUSDTotal, middleware.KeyCostSkipped} + assert.Equal(t, expected, keys, "metadata key allowlist must match the spec") +} + +func TestFactory_AcceptsEmptyAndJSONConfig(t *testing.T) { + configureBuiltin(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) { + configureBuiltin(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 TestFactory_DefaultPricingPathLoadsFixture(t *testing.T) { + configureBuiltin(t) + mw := buildMiddleware(t, nil) + + out, err := mw.Invoke(context.Background(), &middleware.Input{ + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o-mini"}, + {Key: middleware.KeyLLMInputTokens, Value: "1000"}, + {Key: middleware.KeyLLMOutputTokens, Value: "1000"}, + }, + }) + require.NoError(t, err) + require.Equal(t, middleware.DecisionAllow, out.Decision, "cost_meter always allows") + + value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal) + require.True(t, ok, "cost.usd_total must be emitted for known model") + assert.Equal(t, "0.000750", value, "0.00015 + 0.0006 per 1k tokens, 6-decimal format") +} + +func TestFactory_PricingPathOverride(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "custom.yaml"), []byte(fixturePricing), 0o600), "write custom pricing") + builtin.Configure(context.Background(), dir, nil, nil, nil) + + raw, err := json.Marshal(Config{PricingPath: "custom.yaml"}) + require.NoError(t, err) + + mw := buildMiddleware(t, raw) + out, err := mw.Invoke(context.Background(), &middleware.Input{ + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + {Key: middleware.KeyLLMInputTokens, Value: "2000"}, + {Key: middleware.KeyLLMOutputTokens, Value: "1000"}, + }, + }) + require.NoError(t, err) + + value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal) + require.True(t, ok, "cost.usd_total must be emitted with custom pricing path") + assert.Equal(t, "0.015000", value, "2*0.0025 + 1*0.01 = 0.015 with 6-decimal format") +} + +func TestInvoke_ComputesCostForKnownModel(t *testing.T) { + configureBuiltin(t) + mw := buildMiddleware(t, nil) + + out, err := mw.Invoke(context.Background(), &middleware.Input{ + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "anthropic"}, + {Key: middleware.KeyLLMModel, Value: "claude-sonnet-4-5"}, + {Key: middleware.KeyLLMInputTokens, Value: "1000"}, + {Key: middleware.KeyLLMOutputTokens, Value: "1000"}, + }, + }) + require.NoError(t, err) + + value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal) + require.True(t, ok, "cost.usd_total must be emitted") + assert.Equal(t, "0.018000", value, "0.003 + 0.015 = 0.018 with 6-decimal format") + _, skipped := metaValue(t, out.Metadata, middleware.KeyCostSkipped) + assert.False(t, skipped, "cost.skipped must not be set when cost is computed") +} + +func TestInvoke_MissingProvider(t *testing.T) { + configureBuiltin(t) + mw := buildMiddleware(t, nil) + + out, err := mw.Invoke(context.Background(), &middleware.Input{ + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + {Key: middleware.KeyLLMInputTokens, Value: "10"}, + {Key: middleware.KeyLLMOutputTokens, Value: "10"}, + }, + }) + require.NoError(t, err) + value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped) + require.True(t, ok, "cost.skipped must be set when provider is missing") + assert.Equal(t, skipMissingProvider, value, "skip reason matches missing_provider") +} + +func TestInvoke_MissingModel(t *testing.T) { + configureBuiltin(t) + mw := buildMiddleware(t, nil) + + out, err := mw.Invoke(context.Background(), &middleware.Input{ + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMInputTokens, Value: "10"}, + {Key: middleware.KeyLLMOutputTokens, Value: "10"}, + }, + }) + require.NoError(t, err) + value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped) + require.True(t, ok, "cost.skipped must be set when model is missing") + assert.Equal(t, skipMissingModel, value, "skip reason matches missing_model") +} + +func TestInvoke_MissingTokens(t *testing.T) { + configureBuiltin(t) + mw := buildMiddleware(t, nil) + + cases := []struct { + name string + md []middleware.KV + }{ + { + name: "input only", + md: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + {Key: middleware.KeyLLMInputTokens, Value: "10"}, + }, + }, + { + name: "output only", + md: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + {Key: middleware.KeyLLMOutputTokens, Value: "10"}, + }, + }, + { + name: "neither", + md: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + out, err := mw.Invoke(context.Background(), &middleware.Input{Metadata: tc.md}) + require.NoError(t, err) + value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped) + require.True(t, ok, "cost.skipped must be set when token keys are missing") + assert.Equal(t, skipMissingTokens, value, "skip reason matches missing_tokens") + }) + } +} + +func TestInvoke_UnparseableTokens(t *testing.T) { + configureBuiltin(t) + mw := buildMiddleware(t, nil) + + cases := []struct { + name string + in string + out string + }{ + {name: "input non-numeric", in: "abc", out: "10"}, + {name: "output non-numeric", in: "10", out: "xyz"}, + {name: "both garbage", in: "??", out: "??"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + out, err := mw.Invoke(context.Background(), &middleware.Input{ + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + {Key: middleware.KeyLLMInputTokens, Value: tc.in}, + {Key: middleware.KeyLLMOutputTokens, Value: tc.out}, + }, + }) + require.NoError(t, err) + value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped) + require.True(t, ok, "cost.skipped must be set on unparseable tokens") + assert.Equal(t, skipUnparseableTokens, value, "skip reason matches unparseable_tokens") + }) + } +} + +func TestInvoke_ZeroTokens(t *testing.T) { + configureBuiltin(t) + mw := buildMiddleware(t, nil) + + out, err := mw.Invoke(context.Background(), &middleware.Input{ + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + {Key: middleware.KeyLLMInputTokens, Value: "0"}, + {Key: middleware.KeyLLMOutputTokens, Value: "0"}, + }, + }) + require.NoError(t, err) + value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped) + require.True(t, ok, "cost.skipped must be set when both token counts are zero") + assert.Equal(t, skipZeroTokens, value, "skip reason matches zero_tokens") + _, hasCost := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal) + assert.False(t, hasCost, "cost.usd_total must not be emitted for zero tokens") +} + +func TestInvoke_UnknownModel(t *testing.T) { + configureBuiltin(t) + mw := buildMiddleware(t, nil) + + out, err := mw.Invoke(context.Background(), &middleware.Input{ + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "fantasy-model-9000"}, + {Key: middleware.KeyLLMInputTokens, Value: "10"}, + {Key: middleware.KeyLLMOutputTokens, Value: "10"}, + }, + }) + require.NoError(t, err) + value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped) + require.True(t, ok, "cost.skipped must be set when pricing entry is absent") + assert.Equal(t, skipUnknownModel, value, "skip reason matches unknown_model") +} + +func TestInvoke_NilInput(t *testing.T) { + configureBuiltin(t) + mw := buildMiddleware(t, nil) + + out, err := mw.Invoke(context.Background(), nil) + require.NoError(t, err) + require.NotNil(t, out, "output must be returned even on nil input") + assert.Equal(t, middleware.DecisionAllow, out.Decision, "decision must be allow on nil input") + assert.Empty(t, out.Metadata, "no metadata must be emitted on nil input") +} + +const fixturePricingWithCache = `openai: + gpt-4o: + input_per_1k: 0.0025 + output_per_1k: 0.01 + cached_input_per_1k: 0.00125 +anthropic: + claude-sonnet-4-5: + input_per_1k: 0.003 + output_per_1k: 0.015 + cache_read_per_1k: 0.0003 + cache_creation_per_1k: 0.00375 +` + +// configureBuiltinWithCacheRates points the package-level +// FactoryContext at a tmp directory containing pricing entries that +// include the cache rate fields. +func configureBuiltinWithCacheRates(t *testing.T) { + t.Helper() + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "pricing.yaml"), []byte(fixturePricingWithCache), 0o600), "write cache-aware pricing fixture") + builtin.Configure(context.Background(), dir, nil, nil, nil) +} + +// TestInvoke_OpenAICachedSubsetDiscount proves the OpenAI shape end +// to end through the middleware: cached_input_tokens is treated as a +// SUBSET of input_tokens and discounted at the configured rate, not +// added on top. +func TestInvoke_OpenAICachedSubsetDiscount(t *testing.T) { + configureBuiltinWithCacheRates(t) + mw := buildMiddleware(t, nil) + + out, err := mw.Invoke(context.Background(), &middleware.Input{ + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + {Key: middleware.KeyLLMInputTokens, Value: "1000"}, + {Key: middleware.KeyLLMOutputTokens, Value: "500"}, + {Key: middleware.KeyLLMCachedInputTokens, Value: "750"}, + }, + }) + require.NoError(t, err) + require.Equal(t, middleware.DecisionAllow, out.Decision) + + value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal) + require.True(t, ok, "cached subset path must produce a cost — never a skip") + // 250 non-cached at 0.0025/1k + 750 cached at 0.00125/1k + 500 output at 0.01/1k. + assert.Equal(t, "0.006563", value, + "cached subset must be billed at the discount rate, non-cached at the full rate; never double-billed") +} + +// TestInvoke_AnthropicCacheBucketsAdditive proves the Anthropic +// shape: cache_read and cache_creation are additive to input_tokens +// and each carries its own rate. +func TestInvoke_AnthropicCacheBucketsAdditive(t *testing.T) { + configureBuiltinWithCacheRates(t) + mw := buildMiddleware(t, nil) + + out, err := mw.Invoke(context.Background(), &middleware.Input{ + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "anthropic"}, + {Key: middleware.KeyLLMModel, Value: "claude-sonnet-4-5"}, + {Key: middleware.KeyLLMInputTokens, Value: "256"}, + {Key: middleware.KeyLLMOutputTokens, Value: "200"}, + {Key: middleware.KeyLLMCachedInputTokens, Value: "768"}, + {Key: middleware.KeyLLMCacheCreationTokens, Value: "512"}, + }, + }) + require.NoError(t, err) + require.Equal(t, middleware.DecisionAllow, out.Decision) + + value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal) + require.True(t, ok) + // 256 input * 0.003 + 768 cache_read * 0.0003 + 512 cache_creation * 0.00375 + 200 output * 0.015 + // = 0.000768 + 0.0002304 + 0.00192 + 0.003 = 0.0059184 → "0.005918" with 6-decimal format. + assert.Equal(t, "0.005918", value, + "each Anthropic input bucket must bill at its own rate — cache_read cheap, cache_creation expensive, regular input mid") +} + +// TestInvoke_CachedTokensAbsentFallsBackToBaseFormula covers the +// "operator hasn't opted in" path: with no cached metadata keys +// emitted, the meter must produce exactly the same cost as before +// the feature landed. Critical so operators with the new binary but +// no YAML changes see no behavioural drift on OpenAI requests. +func TestInvoke_CachedTokensAbsentFallsBackToBaseFormula(t *testing.T) { + configureBuiltinWithCacheRates(t) + mw := buildMiddleware(t, nil) + + out, err := mw.Invoke(context.Background(), &middleware.Input{ + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + {Key: middleware.KeyLLMInputTokens, Value: "1000"}, + {Key: middleware.KeyLLMOutputTokens, Value: "500"}, + // No KeyLLMCachedInputTokens — the parser didn't see one. + }, + }) + require.NoError(t, err) + value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal) + require.True(t, ok) + // 1000 input * 0.0025 + 500 output * 0.01 = 0.0025 + 0.005 = 0.0075 + assert.Equal(t, "0.007500", value, "no cached metadata = same cost as before the feature landed") +} + +// TestInvoke_UnparseableCachedTokensSkippedSilently proves the +// optional-bucket contract: a malformed cached_input_tokens metadata +// value falls back to 0 (= no cached count) and continues with the +// regular formula. Cache buckets are a refinement, never a reason to +// abort cost computation. +func TestInvoke_UnparseableCachedTokensSkippedSilently(t *testing.T) { + configureBuiltinWithCacheRates(t) + mw := buildMiddleware(t, nil) + + out, err := mw.Invoke(context.Background(), &middleware.Input{ + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + {Key: middleware.KeyLLMInputTokens, Value: "1000"}, + {Key: middleware.KeyLLMOutputTokens, Value: "500"}, + {Key: middleware.KeyLLMCachedInputTokens, Value: "not-a-number"}, + }, + }) + require.NoError(t, err) + value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal) + require.True(t, ok, "garbage cache metadata must NOT switch the response from a cost to a skip — fall back to 0 cached") + assert.Equal(t, "0.007500", value, "same as the no-cached-metadata path") +} + +// TestMiddleware_CloseCancelsReloader proves Close stops the per-instance +// pricing-reload goroutine: a chain rebuild retires the old instance and +// calls Close, which must invoke the cancel func startReloader handed it so +// the mtime-poll loop doesn't outlive the chain. +func TestMiddleware_CloseCancelsReloader(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + m := newMiddleware(nil, cancel) + + require.NoError(t, m.Close(), "Close must not error") + require.Error(t, ctx.Err(), "Close must cancel the reloader context so the poll goroutine exits") +} + +// TestMiddleware_CloseNilSafe confirms Close is a no-op (no panic) for an +// instance with no reloader and for a nil receiver. +func TestMiddleware_CloseNilSafe(t *testing.T) { + require.NoError(t, newMiddleware(nil, nil).Close(), "no-reloader Close must be a no-op") + var m *Middleware + require.NoError(t, m.Close(), "nil-receiver Close must be safe") +} diff --git a/proxy/internal/middleware/builtin/llm_guardrail/factory.go b/proxy/internal/middleware/builtin/llm_guardrail/factory.go new file mode 100644 index 000000000..6dd2a8e8d --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_guardrail/factory.go @@ -0,0 +1,82 @@ +package llm_guardrail + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/proxy/internal/middleware/builtin" +) + +// Config is the JSON-decoded shape accepted by the factory. The +// runtime path consumes the normalised allowlist; raw config is not +// retained beyond construction. +type Config struct { + ModelAllowlist []string `json:"model_allowlist"` + PromptCapture PromptCapture `json:"prompt_capture"` +} + +// PromptCapture toggles the optional prompt capture + redaction step +// that emits llm.request_prompt onto the metadata bag. +type PromptCapture struct { + Enabled bool `json:"enabled"` + RedactPii bool `json:"redact_pii"` +} + +// Factory builds a configured llm_guardrail middleware instance. +type Factory struct{} + +// ID returns the registry identifier matching the middleware ID. +func (Factory) ID() string { return ID } + +// New decodes the raw JSON config and returns a ready Middleware. An +// empty / null / empty-object payload yields a zero-value Config. +func (Factory) New(rawConfig []byte) (middleware.Middleware, error) { + cfg := Config{} + if len(rawConfig) > 0 && !isEmptyJSON(rawConfig) { + if err := json.Unmarshal(rawConfig, &cfg); err != nil { + return nil, fmt.Errorf("decode config: %w", err) + } + } + return New(cfg), nil +} + +// isEmptyJSON reports whether the payload is whitespace, null, or an +// empty object/array. The caller skips Unmarshal in that case so the +// zero-value Config flows through unchanged. +func isEmptyJSON(raw []byte) bool { + trimmed := strings.TrimSpace(string(raw)) + switch trimmed { + case "", "null", "{}", "[]": + return true + } + return false +} + +// normaliseConfig lowercases and trims allowlist entries so the runtime +// match is case-insensitive. Empty entries are dropped. +func normaliseConfig(cfg Config) Config { + if len(cfg.ModelAllowlist) == 0 { + return cfg + } + cleaned := make([]string, 0, len(cfg.ModelAllowlist)) + for _, entry := range cfg.ModelAllowlist { + n := normaliseModel(entry) + if n == "" { + continue + } + cleaned = append(cleaned, n) + } + cfg.ModelAllowlist = cleaned + return cfg +} + +// normaliseModel lowercases and trims a single model identifier. +func normaliseModel(model string) string { + return strings.ToLower(strings.TrimSpace(model)) +} + +func init() { + builtin.Register(Factory{}) +} diff --git a/proxy/internal/middleware/builtin/llm_guardrail/middleware.go b/proxy/internal/middleware/builtin/llm_guardrail/middleware.go new file mode 100644 index 000000000..e6259f06f --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_guardrail/middleware.go @@ -0,0 +1,183 @@ +// Package llm_guardrail implements the SlotOnRequest middleware that +// enforces the per-target LLM guardrail policy: a model allowlist +// check and an opt-in prompt-capture step that may run a PII redactor +// before emitting the prompt into the metadata bag. +// +// The middleware runs after llm_request_parser, which is responsible +// for extracting the model and raw prompt onto the metadata side +// channel. llm_guardrail consumes those keys, decides allow/deny, and +// emits its own decision metadata plus the optional redacted prompt. +package llm_guardrail + +import ( + "context" + "unicode/utf8" + + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +// ID is the registry key for this middleware. +const ID = "llm_guardrail" + +const ( + version = "1.0.0" + maxPromptBytes = 3500 + denyCodeModel = "llm_policy.model_blocked" + denyReasonModel = "model_blocked" + denyMessageModel = "model is not in the policy allowlist" +) + +// Middleware enforces the model allowlist and optionally captures the +// request prompt with PII redaction. +type Middleware struct { + cfg Config +} + +// New constructs a Middleware with the supplied configuration. Model +// allowlist entries are normalised so the runtime check is +// case-insensitive and trim-tolerant. +func New(cfg Config) *Middleware { + return &Middleware{cfg: normaliseConfig(cfg)} +} + +// 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 lists the request body content types the +// middleware needs. Guardrail consumes metadata produced upstream and +// does not touch the body itself, but we keep application/json so the +// body policy retains the parsed payload upstream when required. +func (m *Middleware) AcceptedContentTypes() []string { + return []string{"application/json"} +} + +// MetadataKeys is the closed set of metadata keys this middleware may +// emit. The accumulator drops anything outside this allowlist. +func (m *Middleware) MetadataKeys() []string { + return []string{ + middleware.KeyLLMPolicyDecision, + middleware.KeyLLMPolicyReason, + middleware.KeyLLMRequestPrompt, + } +} + +// MutationsSupported reports whether the middleware emits header / body +// mutations. Guardrail never mutates the request. +func (m *Middleware) MutationsSupported() bool { return false } + +// Invoke runs the policy. The model allowlist is the only deny path; +// prompt capture only affects the metadata emitted alongside an allow. +func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) { + model, modelPresent := lookupMetadata(in.Metadata, middleware.KeyLLMModel) + + if denial := m.evaluateAllowlist(model, modelPresent); denial != nil { + return denial, nil + } + + out := &middleware.Output{ + Decision: middleware.DecisionAllow, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMPolicyDecision, Value: "allow"}, + {Key: middleware.KeyLLMPolicyReason, Value: ""}, + }, + } + + if prompt, ok := m.capturePrompt(in.Metadata); ok { + out.Metadata = append(out.Metadata, middleware.KV{ + Key: middleware.KeyLLMRequestPrompt, + Value: prompt, + }) + } + + return out, nil +} + +// Close releases resources owned by the middleware. Stateless, so this +// is a no-op. +func (m *Middleware) Close() error { return nil } + +// evaluateAllowlist returns a deny Output when the configured allowlist +// rejects the model. A nil return means the request should proceed. +func (m *Middleware) evaluateAllowlist(model string, modelPresent bool) *middleware.Output { + if len(m.cfg.ModelAllowlist) == 0 { + return nil + } + if !modelPresent { + return nil + } + if m.modelInAllowlist(model) { + return nil + } + return &middleware.Output{ + Decision: middleware.DecisionDeny, + DenyStatus: 403, + DenyReason: &middleware.DenyReason{ + Code: denyCodeModel, + Message: denyMessageModel, + Details: map[string]string{"model": model}, + }, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMPolicyDecision, Value: "deny"}, + {Key: middleware.KeyLLMPolicyReason, Value: denyReasonModel}, + }, + } +} + +// modelInAllowlist reports whether the model matches any allowlist +// entry under the case-insensitive, trim-tolerant comparison rule. +func (m *Middleware) modelInAllowlist(model string) bool { + normalised := normaliseModel(model) + if normalised == "" { + return false + } + for _, allowed := range m.cfg.ModelAllowlist { + if allowed == normalised { + return true + } + } + return false +} + +// capturePrompt returns the prompt to emit and whether it should be +// emitted at all. The truncation guarantee is upheld here regardless of +// whether redaction grew the string. +func (m *Middleware) capturePrompt(meta []middleware.KV) (string, bool) { + if !m.cfg.PromptCapture.Enabled { + return "", false + } + raw, ok := lookupMetadata(meta, middleware.KeyLLMRequestPromptRaw) + if !ok { + return "", false + } + prompt := raw + if m.cfg.PromptCapture.RedactPii { + prompt = redactPII(prompt) + } + if len(prompt) > maxPromptBytes { + // Back off to a UTF-8 rune boundary so we never emit a string + // split mid-rune. + cut := maxPromptBytes + for cut > 0 && !utf8.RuneStart(prompt[cut]) { + cut-- + } + prompt = prompt[:cut] + } + return prompt, true +} + +// lookupMetadata finds the first KV with the given key. Returns the +// value and true when present; the empty string and false otherwise. +func lookupMetadata(meta []middleware.KV, key string) (string, bool) { + for _, kv := range meta { + if kv.Key == key { + return kv.Value, true + } + } + return "", false +} diff --git a/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go b/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go new file mode 100644 index 000000000..865dc07af --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go @@ -0,0 +1,219 @@ +package llm_guardrail + +import ( + "context" + "encoding/json" + "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 newInput(meta ...middleware.KV) *middleware.Input { + return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: meta} +} + +func TestMiddlewareIdentity(t *testing.T) { + mw := New(Config{}) + assert.Equal(t, ID, mw.ID(), "middleware ID must be llm_guardrail") + assert.Equal(t, "1.0.0", mw.Version(), "version must be 1.0.0") + assert.Equal(t, middleware.SlotOnRequest, mw.Slot(), "guardrail must run in SlotOnRequest") + assert.False(t, mw.MutationsSupported(), "guardrail must not mutate requests") + assert.Equal(t, []string{"application/json"}, mw.AcceptedContentTypes(), "guardrail accepts application/json bodies") + assert.Equal(t, + []string{ + middleware.KeyLLMPolicyDecision, + middleware.KeyLLMPolicyReason, + middleware.KeyLLMRequestPrompt, + }, + mw.MetadataKeys(), + "metadata key allowlist must match the spec", + ) + require.NoError(t, mw.Close()) +} + +func TestAllowlistEmptyAllowsAnyModel(t *testing.T) { + mw := New(Config{}) + out, err := mw.Invoke(context.Background(), newInput( + middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + )) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "empty allowlist must allow any model") + v, ok := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision) + require.True(t, ok, "decision metadata must be emitted") + assert.Equal(t, "allow", v, "decision must be allow") + r, ok := metaValue(t, out.Metadata, middleware.KeyLLMPolicyReason) + require.True(t, ok, "reason metadata must be emitted") + assert.Equal(t, "", r, "reason must be empty on allow") +} + +func TestAllowlistMatchAllows(t *testing.T) { + mw := New(Config{ModelAllowlist: []string{"gpt-4o", "claude-opus-4"}}) + out, err := mw.Invoke(context.Background(), newInput( + middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + )) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "model in allowlist must be allowed") +} + +func TestAllowlistMissDenies(t *testing.T) { + mw := New(Config{ModelAllowlist: []string{"gpt-4o"}}) + out, err := mw.Invoke(context.Background(), newInput( + middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"}, + )) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionDeny, out.Decision, "non-allowlisted model must be denied") + assert.Equal(t, 403, out.DenyStatus, "deny status must be 403") + require.NotNil(t, out.DenyReason, "deny reason must be populated") + assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must match spec") + assert.Equal(t, "model is not in the policy allowlist", out.DenyReason.Message, "deny message must match spec") + assert.Equal(t, "claude-opus-4", out.DenyReason.Details["model"], "deny details must include the offending model") + + dec, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision) + assert.Equal(t, "deny", dec, "decision metadata must be deny") + reason, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyReason) + assert.Equal(t, "model_blocked", reason, "reason metadata must be model_blocked") +} + +func TestAllowlistCaseInsensitive(t *testing.T) { + mw := New(Config{ModelAllowlist: []string{" GPT-4o ", "Claude-OPUS-4"}}) + cases := []string{"gpt-4o", "GPT-4O", " claude-opus-4 "} + for _, model := range cases { + out, err := mw.Invoke(context.Background(), newInput( + middleware.KV{Key: middleware.KeyLLMModel, Value: model}, + )) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "case/whitespace variants must match: %q", model) + } +} + +func TestAllowlistMissingModelKeyAllows(t *testing.T) { + mw := New(Config{ModelAllowlist: []string{"gpt-4o"}}) + out, err := mw.Invoke(context.Background(), newInput()) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "missing model key must allow even with non-empty allowlist") + dec, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision) + assert.Equal(t, "allow", dec, "decision must be allow when model key is absent") +} + +func TestPromptCaptureDisabledEmitsNoPrompt(t *testing.T) { + mw := New(Config{}) + out, err := mw.Invoke(context.Background(), newInput( + middleware.KV{Key: middleware.KeyLLMRequestPromptRaw, Value: "hello world"}, + )) + require.NoError(t, err) + _, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPrompt) + assert.False(t, ok, "prompt must not be emitted when capture is disabled") +} + +func TestPromptCaptureNoRedactionEmitsRaw(t *testing.T) { + mw := New(Config{PromptCapture: PromptCapture{Enabled: true}}) + raw := "hello world from user@example.com" + out, err := mw.Invoke(context.Background(), newInput( + middleware.KV{Key: middleware.KeyLLMRequestPromptRaw, Value: raw}, + )) + require.NoError(t, err) + prompt, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPrompt) + require.True(t, ok, "prompt must be emitted when capture is enabled") + assert.Equal(t, raw, prompt, "prompt must pass through unchanged when redaction is off") +} + +func TestPromptCaptureWithRedactionRedacts(t *testing.T) { + mw := New(Config{PromptCapture: PromptCapture{Enabled: true, RedactPii: true}}) + raw := "contact me at user@example.com or +14155551234" + out, err := mw.Invoke(context.Background(), newInput( + middleware.KV{Key: middleware.KeyLLMRequestPromptRaw, Value: raw}, + )) + require.NoError(t, err) + prompt, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPrompt) + require.True(t, ok, "prompt must be emitted when capture is enabled") + assert.Contains(t, prompt, "[REDACTED:email]", "email must be redacted") + assert.Contains(t, prompt, "[REDACTED:phone]", "phone must be redacted") + assert.NotContains(t, prompt, "user@example.com", "raw email must not leak") +} + +func TestPromptCaptureRedactionTruncatesIfGrows(t *testing.T) { + mw := New(Config{PromptCapture: PromptCapture{Enabled: true, RedactPii: true}}) + body := strings.Repeat("a", maxPromptBytes-10) + " user@example.com" + out, err := mw.Invoke(context.Background(), newInput( + middleware.KV{Key: middleware.KeyLLMRequestPromptRaw, Value: body}, + )) + require.NoError(t, err) + prompt, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPrompt) + require.True(t, ok, "prompt must be emitted when capture is enabled") + assert.LessOrEqual(t, len(prompt), maxPromptBytes, "prompt must be truncated to maxPromptBytes") +} + +func TestPromptCaptureMissingRawNoEmit(t *testing.T) { + mw := New(Config{PromptCapture: PromptCapture{Enabled: true, RedactPii: true}}) + out, err := mw.Invoke(context.Background(), newInput()) + require.NoError(t, err) + _, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPrompt) + assert.False(t, ok, "prompt must not be emitted when raw key is missing") +} + +func TestFactoryAcceptsZeroConfigs(t *testing.T) { + cases := map[string][]byte{ + "nil": nil, + "empty": []byte(""), + "whitespace": []byte(" \n "), + "null": []byte("null"), + "emptyObject": []byte("{}"), + } + f := Factory{} + for name, raw := range cases { + mw, err := f.New(raw) + require.NoError(t, err, "case %s must yield a zero-value config", name) + require.NotNil(t, mw) + assert.Equal(t, ID, mw.ID(), "case %s must build a guardrail middleware", name) + } +} + +func TestFactoryDecodesValidConfig(t *testing.T) { + cfg := Config{ + ModelAllowlist: []string{"gpt-4o"}, + PromptCapture: PromptCapture{Enabled: true, RedactPii: true}, + } + raw, err := json.Marshal(cfg) + require.NoError(t, err, "marshalling test config must succeed") + mw, err := Factory{}.New(raw) + require.NoError(t, err) + require.NotNil(t, mw) +} + +func TestFactoryRejectsMalformedJSON(t *testing.T) { + mw, err := Factory{}.New([]byte("{not-json")) + assert.Error(t, err, "malformed JSON must surface as a factory error") + assert.Nil(t, mw, "no middleware must be returned on malformed config") +} + +func TestFactoryNormalisesAllowlist(t *testing.T) { + raw := []byte(`{"model_allowlist":[" GPT-4o ","",""," Claude-3 "]}`) + mw, err := Factory{}.New(raw) + require.NoError(t, err) + out, err := mw.Invoke(context.Background(), newInput( + middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + )) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "factory must lowercase + trim allowlist entries") + out2, err := mw.Invoke(context.Background(), newInput( + middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-3"}, + )) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out2.Decision, "trimmed entry must still match") +} diff --git a/proxy/internal/middleware/builtin/llm_guardrail/redact.go b/proxy/internal/middleware/builtin/llm_guardrail/redact.go new file mode 100644 index 000000000..c6cb270df --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_guardrail/redact.go @@ -0,0 +1,75 @@ +package llm_guardrail + +import ( + "regexp" + "strings" + + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +// PII redactor scope: redact prompt content BEFORE it lands in the metadata +// bag. The bearer-with-keyword pass runs first so the keyword is preserved. +// We then chain the package-level middleware.Scan to pick up PEM, JWT, AWS +// access keys, generic bearer tokens (40+ chars), and Luhn-validated credit +// cards — keeping prompt redaction in sync with metadata-value scanning. Email, +// SSN (dashed form), phone (E.164 + NA), and IPv4 are prompt-shaped patterns +// the metadata scanner intentionally leaves alone. +var ( + emailRegex = regexp.MustCompile(`[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}`) + ssnRegex = regexp.MustCompile(`\b\d{3}-\d{2}-\d{4}\b`) + phoneE164 = regexp.MustCompile(`\+\d{8,15}\b`) + // phoneNARgx accepts the 3-3-4 North-American shape with any of the common + // separators (space, dot, dash, slash) or none at all between the area code + // and the body. The optional `\(?...\)?` wraps the area code; the separator + // classes use `*` (not `?`) so multi-char separators ("(202) " followed by + // space-and-something) and zero-separator runs ("2025550134") both match. + // False-positive tradeoff: 10 consecutive digits in a prompt will be + // treated as a phone number. For PII redaction that is the correct way to + // err — under-redaction leaks; over-redaction is annoying. + phoneNARgx = regexp.MustCompile(`\(?\b\d{3}\)?[\s.\-/]*\d{3}[\s.\-/]*\d{4}\b`) + bearerRegex = regexp.MustCompile(`(?i)\b(bearer|token|api[_-]?key|authorization)([\s:=]+)(\S{20,})`) + ipv4Regex = regexp.MustCompile(`\b(?:(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)\.){3}(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)\b`) +) + +// redactPII is the package-private alias kept for internal callers; new code +// outside the guardrail middleware should call RedactPII. +func redactPII(value string) string { return RedactPII(value) } + +// RedactPII replaces high-signal PII patterns in value with +// `[REDACTED:]`. Non-matching input is returned unchanged. Exported so +// the request / response parsers can reuse the same coverage on raw prompts +// and completions when the account's redact_pii toggle is on. +func RedactPII(value string) string { + if value == "" { + return value + } + result := value + // Keyword-preserving bearer first so the "bearer "/"token=" prefix survives + // before the generic scanner gets at the same content. + result = bearerRegex.ReplaceAllStringFunc(result, redactBearer) + // Structured secrets shared with metadata-value scanning: PEM, JWT, AWS + // keys, generic bearer (40+), and Luhn-validated credit cards. + result = middleware.Scan(result) + // Prompt-shaped PII the metadata scanner doesn't cover. + result = emailRegex.ReplaceAllString(result, "[REDACTED:email]") + result = ssnRegex.ReplaceAllString(result, "[REDACTED:ssn]") + result = phoneE164.ReplaceAllString(result, "[REDACTED:phone]") + result = phoneNARgx.ReplaceAllString(result, "[REDACTED:phone]") + result = ipv4Regex.ReplaceAllString(result, "[REDACTED:ip]") + return result +} + +// redactBearer keeps the leading keyword and its separator, replacing +// only the secret payload so the surrounding context is preserved. +func redactBearer(match string) string { + sub := bearerRegex.FindStringSubmatch(match) + if len(sub) < 4 { + return "[REDACTED:bearer]" + } + var b strings.Builder + b.Grow(len(sub[1]) + len(sub[2]) + len("[REDACTED:bearer]")) + b.WriteString(sub[1]) + b.WriteString(sub[2]) + b.WriteString("[REDACTED:bearer]") + return b.String() +} diff --git a/proxy/internal/middleware/builtin/llm_guardrail/redact_test.go b/proxy/internal/middleware/builtin/llm_guardrail/redact_test.go new file mode 100644 index 000000000..ef17f1d0a --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_guardrail/redact_test.go @@ -0,0 +1,217 @@ +package llm_guardrail + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestRedactPIIEmptyInput(t *testing.T) { + assert.Equal(t, "", redactPII(""), "empty input must round-trip unchanged") +} + +func TestRedactPIIPlainTextUntouched(t *testing.T) { + in := "the quick brown fox jumps over the lazy dog" + assert.Equal(t, in, redactPII(in), "non-PII text must pass through unchanged") +} + +func TestRedactPIIEmail(t *testing.T) { + cases := []string{ + "contact user@example.com today", + "first.last+tag@sub.example.co", + "USER_42@EXAMPLE.COM", + } + for _, in := range cases { + out := redactPII(in) + assert.Contains(t, out, "[REDACTED:email]", "email must be redacted in %q", in) + assert.NotContains(t, strings.ToLower(out), "@example", "raw email host must not survive in %q", in) + } +} + +func TestRedactPIISSN(t *testing.T) { + in := "ssn 123-45-6789 should be hidden" + out := redactPII(in) + assert.Contains(t, out, "[REDACTED:ssn]", "SSN must be redacted") + assert.NotContains(t, out, "123-45-6789", "raw SSN must not survive") +} + +func TestRedactPIIPhoneE164(t *testing.T) { + in := "call me at +14155551234 anytime" + out := redactPII(in) + assert.Contains(t, out, "[REDACTED:phone]", "E.164 phone must be redacted") + assert.NotContains(t, out, "+14155551234", "raw E.164 phone must not survive") +} + +func TestRedactPIIPhoneNorthAmerican(t *testing.T) { + cases := []string{ + "call (415) 555-1234 now", + "call 415-555-1234 now", + "call 415.555.1234 now", + "call 415 555 1234 now", + } + for _, in := range cases { + out := redactPII(in) + assert.Contains(t, out, "[REDACTED:phone]", "NA phone must be redacted in %q", in) + assert.NotContains(t, out, "555-1234", "raw NA phone must not survive in %q", in) + } +} + +func TestRedactPIIBearerKeepsKeyword(t *testing.T) { + cases := []struct { + in string + keyword string + }{ + {"Authorization: Bearer abcdefghijklmnopqrstuvwxyz0123", "Bearer"}, + {"token = abcdefghijklmnopqrstuvwxyz", "token"}, + {"api_key=abcdefghijklmnopqrstuvwxyz0123", "api_key"}, + {"API-KEY: abcdefghijklmnopqrstuvwxyz0123", "API-KEY"}, + {"authorization: abcdefghijklmnopqrstuvwxyz0123", "authorization"}, + } + for _, tc := range cases { + out := redactPII(tc.in) + assert.Contains(t, out, "[REDACTED:bearer]", "bearer-style secret must be redacted in %q", tc.in) + assert.Contains(t, out, tc.keyword, "leading keyword %q must be preserved in %q", tc.keyword, tc.in) + assert.NotContains(t, out, "abcdefghijklmnopqrstuvwxyz0123", "raw bearer payload must not survive in %q", tc.in) + } +} + +func TestRedactPIIBearerShortValueUntouched(t *testing.T) { + in := "token=short" + out := redactPII(in) + assert.Equal(t, in, out, "short bearer-style values must not be redacted") +} + +func TestRedactPIICombined(t *testing.T) { + in := "email user@example.com phone +14155551234 ssn 123-45-6789 token abcdefghijklmnopqrstuvwxyz0123" + out := redactPII(in) + assert.Contains(t, out, "[REDACTED:email]", "email must be redacted in combined input") + assert.Contains(t, out, "[REDACTED:phone]", "phone must be redacted in combined input") + assert.Contains(t, out, "[REDACTED:ssn]", "SSN must be redacted in combined input") + assert.Contains(t, out, "[REDACTED:bearer]", "bearer must be redacted in combined input") + assert.NotContains(t, out, "user@example.com", "raw email must not survive combined input") + assert.NotContains(t, out, "+14155551234", "raw phone must not survive combined input") + assert.NotContains(t, out, "123-45-6789", "raw SSN must not survive combined input") +} + +func TestRedactPIICreditCard(t *testing.T) { + // 4242424242424242 is a well-known Stripe test number (Visa, Luhn-valid). + cases := []string{ + "please charge 4242424242424242 now", + "card: 4242-4242-4242-4242", + "4242 4242 4242 4242 expires 12/30", + } + for _, in := range cases { + out := redactPII(in) + assert.Contains(t, out, "[REDACTED:cc]", "Luhn-valid credit card must be redacted in %q", in) + assert.NotContains(t, out, "4242424242424242", "raw card digits must not survive in %q", in) + assert.NotContains(t, out, "4242-4242-4242-4242", "raw dashed card must not survive in %q", in) + } +} + +func TestRedactPIIIPv4(t *testing.T) { + cases := []string{ + "connect to 10.0.42.7 over the tunnel", + "server 192.168.1.100 down", + "public address 203.0.113.42 was hit", + } + for _, in := range cases { + out := redactPII(in) + assert.Contains(t, out, "[REDACTED:ip]", "IPv4 must be redacted in %q", in) + } +} + +func TestRedactPIIJWT(t *testing.T) { + // No "token "/"bearer " prefix here, so the bearer-with-keyword pass leaves + // it alone and the JWT pattern from middleware.Scan must catch it. + in := "session eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJ1c2VyXzQyIn0.signaturepart expires soon" + out := redactPII(in) + assert.Contains(t, out, "[REDACTED:jwt]", "JWT must be redacted when no bearer keyword precedes it") + assert.NotContains(t, out, "eyJhbGciOiJIUzI1NiJ9", "raw JWT header must not survive") +} + +func TestRedactPIIAWSAccessKey(t *testing.T) { + in := "the key AKIAIOSFODNN7EXAMPLE belongs to test user" + out := redactPII(in) + assert.Contains(t, out, "[REDACTED:aws_key]", "AWS access key must be redacted") + assert.NotContains(t, out, "AKIAIOSFODNN7EXAMPLE", "raw AWS key must not survive") +} + +func TestRedactPIIPlainNumbersUntouched(t *testing.T) { + // 1234567890123 is 13 digits but fails Luhn; must NOT trip the CC redactor. + // We use a 13-digit value (the CC-candidate range starts at 13) so the only + // risk is the CC pattern firing. Phone redaction is 10-digit by design and + // would catch 1234567890123 as a phone — that's expected and not what this + // test guards against. + in := "order number 1234567890123 is queued" + out := redactPII(in) + assert.NotContains(t, out, "[REDACTED:cc]", "non-Luhn digit sequences must not be redacted as credit cards") +} + +// piiFixture mirrors the user-supplied test fixture: each record carries one +// email, one SSN, and one phone in a representative format. The test asserts +// that EVERY raw token disappears after redaction and the right [REDACTED:*] +// markers show up. Names are kept in the input and must survive — names are +// not a pattern the redactor tries to catch. +type piiFixture struct { + name string // person name (must survive redaction) + email string + ssn string + phone string +} + +var fixtureRecords = []piiFixture{ + {"Alice Johnson", "alice.johnson@example.com", "123-45-6789", "(202) 555-0147"}, + {"Brian Smith", "brian.smith@example.org", "987-65-4321", "202-555-0163"}, + {"Carla Nguyen", "c.nguyen@test.local", "111-22-3333", "+1-202-555-0188"}, + {"David Martinez", "david.martinez@example.com", "222-33-4444", "202.555.0199"}, + {"Evelyn Parker", "evelyn.parker@example.org", "333-44-5555", "1-202-555-0112"}, + {"Frank O'Connor", "frank.oconnor@test.local", "444-55-6666", "2025550134"}, + {"Grace Lee", "grace.lee@example.com", "555-66-7777", "(202)555-0156"}, + {"Hassan Ali", "hassan.ali@example.org", "666-77-8888", "+1 (202) 555-0175"}, + {"Isabella Rossi", "i.rossi@test.local", "777-88-9999", "202 555 0121"}, + {"Jamal Thompson", "jamal.thompson@example.com", "888-99-0001", "202/555/0108"}, +} + +// TestRedactPII_FixtureRecord drives every record through redactPII and +// asserts the email, SSN, and phone are all redacted, the name survives, and +// the appropriate REDACTED markers are present. This is the spec the redactor +// must meet for the kind of prompts operators throw at it. +func TestRedactPII_FixtureRecord(t *testing.T) { + for _, rec := range fixtureRecords { + t.Run(rec.name, func(t *testing.T) { + in := "Name: " + rec.name + "\n Email: " + rec.email + "\n SSN: " + rec.ssn + "\n Phone: " + rec.phone + out := redactPII(in) + + assert.Contains(t, out, rec.name, "name must survive (not a PII pattern the redactor catches)") + assert.Contains(t, out, "[REDACTED:email]", "email marker must appear for %q", rec.email) + assert.Contains(t, out, "[REDACTED:ssn]", "ssn marker must appear for %q", rec.ssn) + assert.Contains(t, out, "[REDACTED:phone]", "phone marker must appear for %q", rec.phone) + + assert.NotContains(t, out, rec.email, "raw email must not survive: %q", rec.email) + assert.NotContains(t, out, rec.ssn, "raw SSN must not survive: %q", rec.ssn) + // Phone: assert the local digits (last 7) are gone. Country-code + // remnants like "+1 " or "1-" may remain in front of the redaction + // because the E.164 pattern needs digits-only after '+' — that's + // acceptable, the personally-identifying portion is removed. + localDigits := lastSevenDigits(rec.phone) + assert.NotContains(t, out, localDigits, "raw phone local digits %q must not survive in redacted output of %q", localDigits, rec.phone) + }) + } +} + +// lastSevenDigits returns the last 7 digits of a phone number, ignoring +// formatting. It's the unique "subscriber" portion that absolutely must be +// scrubbed regardless of which prefix the redactor leaves behind. +func lastSevenDigits(phone string) string { + digits := make([]byte, 0, len(phone)) + for i := 0; i < len(phone); i++ { + if phone[i] >= '0' && phone[i] <= '9' { + digits = append(digits, phone[i]) + } + } + if len(digits) <= 7 { + return string(digits) + } + return string(digits[len(digits)-7:]) +} diff --git a/proxy/internal/middleware/builtin/llm_identity_inject/factory.go b/proxy/internal/middleware/builtin/llm_identity_inject/factory.go new file mode 100644 index 000000000..8594c392d --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_identity_inject/factory.go @@ -0,0 +1,108 @@ +package llm_identity_inject + +import ( + "bytes" + "encoding/json" + "fmt" + "strings" + + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/proxy/internal/middleware/builtin" +) + +// ProviderInjection describes one resolved provider's injection rule. +// Identity stamping uses one of HeaderPair / JSONMetadata; ExtraHeaders +// is independent — each entry is a static (operator-configured) header +// stamped on every matching request with anti-spoof. A rule with no +// shape AND no extras is dropped at New() time as a no-op. +type ProviderInjection struct { + // ProviderID is the resolved provider id — matches the value + // llm_router stamps under KeyLLMResolvedProviderID. + ProviderID string `json:"provider_id"` + // HeaderPair is the LiteLLM-style wire convention: separate + // headers for end-user id and tags CSV. + HeaderPair *HeaderPairRule `json:"header_pair,omitempty"` + // JSONMetadata is the Portkey-style wire convention: a single + // header carrying a JSON object keyed by reserved field names. + JSONMetadata *JSONMetadataRule `json:"json_metadata,omitempty"` + // ExtraHeaders is an operator-configured list of static headers + // (e.g. "x-portkey-config: pc-...") that the middleware stamps + // on every matching request. The synth pre-resolves the values + // from the provider record's ExtraValues map; the middleware + // just emits them. Each name is also added to HeadersRemove for + // anti-spoof so a client can't smuggle their own value. + ExtraHeaders []ExtraHeaderKV `json:"extra_headers,omitempty"` +} + +// ExtraHeaderKV is one static header entry the middleware stamps as-is. +type ExtraHeaderKV struct { + Name string `json:"name"` + Value string `json:"value"` +} + +// HeaderPairRule emits identity through dedicated per-dimension +// headers. The two *InBody flags layer body-level identity on top: when +// TagsInBody is set the middleware also writes the tag list into the +// request body's metadata.tags array (required for LiteLLM tag-budget +// enforcement, which only inspects the body); when EndUserIDInBody is +// set the display identity is also written into the body's top-level +// "user" field (the OpenAI-standard end-user identifier — defense-in- +// depth and anti-spoof on top of the header path). +type HeaderPairRule struct { + EndUserIDHeader string `json:"end_user_id_header,omitempty"` + TagsHeader string `json:"tags_header,omitempty"` + TagsInBody bool `json:"tags_in_body,omitempty"` + EndUserIDInBody bool `json:"end_user_id_in_body,omitempty"` +} + +// JSONMetadataRule emits identity through a single JSON-object header. +// Empty UserKey/GroupsKey skip that dimension at emit time. When +// MaxValueLength > 0 each emitted JSON value is truncated to that many +// bytes — Portkey enforces 128 chars per value. +type JSONMetadataRule struct { + Header string `json:"header"` + UserKey string `json:"user_key,omitempty"` + GroupsKey string `json:"groups_key,omitempty"` + MaxValueLength int `json:"max_value_length,omitempty"` +} + +// Config is the on-wire configuration accepted by the factory. An +// empty Providers slice yields a no-op middleware (every resolved +// provider passes through unchanged). +type Config struct { + Providers []ProviderInjection `json:"providers"` +} + +// Factory builds llm_identity_inject 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 +// yield a no-op middleware. Non-empty payloads must parse cleanly so +// misconfigurations surface at chain build time. +func (Factory) New(rawConfig []byte) (middleware.Middleware, error) { + cfg := Config{} + if !isEmptyJSON(rawConfig) { + if err := json.Unmarshal(rawConfig, &cfg); err != nil { + return nil, fmt.Errorf("decode config: %w", err) + } + } + return New(cfg), nil +} + +// isEmptyJSON reports whether the payload is whitespace, null, or an +// empty object/array. +func isEmptyJSON(raw []byte) bool { + trimmed := strings.TrimSpace(string(bytes.TrimSpace(raw))) + switch trimmed { + case "", "null", "{}", "[]": + return true + } + return false +} + +func init() { + builtin.Register(Factory{}) +} diff --git a/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go b/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go new file mode 100644 index 000000000..96f686c0e --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go @@ -0,0 +1,440 @@ +// Package llm_identity_inject implements the SlotOnRequest middleware +// that stamps the caller's NetBird identity onto upstream LLM-gateway +// requests. It runs after llm_router (which resolves the provider) and +// looks up the resolved provider id against a per-account injection +// table built by the synthesiser from the catalog's IdentityInjection +// metadata. +// +// Two wire shapes are supported, dispatched per-rule: +// +// - HeaderPair (LiteLLM-style): separate end-user-id and tags +// headers; tags emitted as a CSV value. +// - JSONMetadata (Portkey-style): one header carrying a JSON +// object with reserved keys for user / groups; per-value byte +// length capped when the rule sets MaxValueLength. +// +// In both cases, identity comes from Input.UserEmail (peer-attached +// user's email or peer.Name fallback) and groups come from the +// authorising-groups intersection llm_router emitted (with +// id→display-name translation via Input.UserGroups / UserGroupNames +// positional pairing). HeadersRemove runs before HeadersAdd in the +// framework, so a client can never spoof identity by stamping these +// headers themselves. +package llm_identity_inject + +import ( + "context" + "encoding/json" + "sort" + "strings" + + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +// ID is the registry identifier for this middleware. +const ID = "llm_identity_inject" + +// Version is reported via Middleware.Version(). +const Version = "1.0.0" + +// Middleware stamps NetBird identity onto upstream requests for the +// configured set of resolved providers. +type Middleware struct { + cfg Config + byID map[string]ProviderInjection +} + +// New constructs a Middleware from the supplied configuration. A nil +// or empty Providers slice yields a no-op middleware. +func New(cfg Config) *Middleware { + byID := make(map[string]ProviderInjection, len(cfg.Providers)) + for _, p := range cfg.Providers { + if p.ProviderID == "" { + continue + } + // Drop entries that wouldn't inject anything — keeps the + // runtime check tight. Also drop entries that set both + // shapes (configuration error; refuse to guess which wins). + // Extras alone are enough to keep the rule alive even if + // neither identity shape is set. + hasExtras := false + for _, e := range p.ExtraHeaders { + if e.Name != "" && e.Value != "" { + hasExtras = true + break + } + } + switch { + case p.HeaderPair != nil && p.JSONMetadata != nil: + continue + case p.HeaderPair != nil: + if p.HeaderPair.EndUserIDHeader == "" && p.HeaderPair.TagsHeader == "" && !p.HeaderPair.TagsInBody && !p.HeaderPair.EndUserIDInBody && !hasExtras { + continue + } + case p.JSONMetadata != nil: + if p.JSONMetadata.Header == "" { + continue + } + if p.JSONMetadata.UserKey == "" && p.JSONMetadata.GroupsKey == "" && !hasExtras { + continue + } + default: + if !hasExtras { + continue + } + } + byID[p.ProviderID] = p + } + return &Middleware{cfg: cfg, byID: byID} +} + +// 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 — this middleware reads only +// metadata and identity fields on the Input envelope. +func (m *Middleware) AcceptedContentTypes() []string { return nil } + +// MetadataKeys is empty: the middleware emits no metadata. Identity +// stamping is a header-only operation. +func (m *Middleware) MetadataKeys() []string { return nil } + +// MutationsSupported reports that the middleware emits header +// mutations on the Output envelope. +func (m *Middleware) MutationsSupported() bool { return true } + +// Close releases resources owned by the middleware. Stateless, so +// this is a no-op. +func (m *Middleware) Close() error { return nil } + +// Invoke stamps identity headers when the resolved provider has an +// injection rule. Always Allow. +func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) { + out := &middleware.Output{Decision: middleware.DecisionAllow} + if len(m.byID) == 0 || in == nil { + return out, nil + } + resolved, ok := lookupMetadata(in.Metadata, middleware.KeyLLMResolvedProviderID) + if !ok || resolved == "" { + return out, nil + } + rule, ok := m.byID[resolved] + if !ok { + return out, nil + } + + var mutations *middleware.Mutations + switch { + case rule.HeaderPair != nil: + mutations = applyHeaderPair(rule.HeaderPair, in) + case rule.JSONMetadata != nil: + mutations = applyJSONMetadata(rule.JSONMetadata, in) + } + + // ExtraHeaders are independent of the identity shape. Stamp each + // non-empty entry with anti-spoof: Remove first (frame strips it + // before our Add lands) so a client can't smuggle a value, then + // Add our trusted one. + if len(rule.ExtraHeaders) > 0 { + if mutations == nil { + mutations = &middleware.Mutations{} + } + for _, h := range rule.ExtraHeaders { + if h.Name == "" || h.Value == "" { + continue + } + mutations.HeadersRemove = append(mutations.HeadersRemove, h.Name) + mutations.HeadersAdd = append(mutations.HeadersAdd, middleware.KV{ + Key: h.Name, + Value: h.Value, + }) + } + } + + if mutations == nil || (len(mutations.HeadersAdd) == 0 && len(mutations.HeadersRemove) == 0 && len(mutations.BodyReplace) == 0) { + return out, nil + } + out.Mutations = mutations + return out, nil +} + +// applyHeaderPair builds the LiteLLM-style mutations: separate per- +// dimension headers, with anti-spoof Removes paired with trusted Adds. +func applyHeaderPair(rule *HeaderPairRule, in *middleware.Input) *middleware.Mutations { + mutations := &middleware.Mutations{} + + if rule.EndUserIDHeader != "" { + mutations.HeadersRemove = append(mutations.HeadersRemove, rule.EndUserIDHeader) + // Prefer the email when the auth path carried it: gateways + // like LiteLLM key per-user budgets and dashboards on a + // human-readable identifier; the user_id is an opaque + // management-server primary key. Fall back to user_id when + // no email is available (non-OIDC schemes, legacy JWTs). + if identity := identityFor(in); identity != "" { + mutations.HeadersAdd = append(mutations.HeadersAdd, middleware.KV{ + Key: rule.EndUserIDHeader, + Value: identity, + }) + } + } + + if rule.TagsHeader != "" { + mutations.HeadersRemove = append(mutations.HeadersRemove, rule.TagsHeader) + if csv := authorisingTagsCSV(in); csv != "" { + mutations.HeadersAdd = append(mutations.HeadersAdd, middleware.KV{ + Key: rule.TagsHeader, + Value: csv, + }) + } + } + + if rule.TagsInBody || rule.EndUserIDInBody { + // Body-level identity unlocks gateway behaviour the header + // path can't reach (LiteLLM's _tag_max_budget_check only + // inspects the body; OpenAI direct only reads the body's + // "user" field for attribution). The header path stays + // intact, so we still get attribution + per-end-user budget + // gating when body inject can't run (truncated body, + // non-JSON, hostile metadata shape). + var bodyTags []string + if rule.TagsInBody { + bodyTags = authorisingTagsSlice(in) + } + var bodyUser string + if rule.EndUserIDInBody { + bodyUser = identityFor(in) + } + if newBody, ok := injectIntoBody(in, bodyTags, bodyUser); ok { + mutations.BodyReplace = newBody + } + } + + return mutations +} + +// injectIntoBody parses the request body and writes the supplied +// identity dimensions into it. Tags land at metadata.tags (creating +// the metadata object when absent); the user identity lands at the +// top-level "user" field (OpenAI-standard end-user identifier). +// Returns the re-marshaled body and ok=true when at least one field +// was written. Returns ok=false (no mutation) when: +// +// - both inputs are empty (nothing to write); +// - the body is empty or truncated (we don't have the full document +// to safely round-trip); +// - the body isn't a JSON object (skip silently — this middleware +// only knows how to inject into OpenAI-compatible JSON payloads). +// +// A non-object existing `metadata` field skips the tag write but +// still allows the user write to land — we don't clobber the client's +// non-object metadata, but the orthogonal user field is fair game. +// The header path emission still runs in skip cases, so spend tracking +// + header-resolved end-user budgets continue to work without body- +// level enforcement. +func injectIntoBody(in *middleware.Input, tags []string, userID string) ([]byte, bool) { + wantTags := len(tags) > 0 + wantUser := userID != "" + if !wantTags && !wantUser { + return nil, false + } + if in == nil || len(in.Body) == 0 || in.BodyTruncated { + return nil, false + } + var doc map[string]any + if err := json.Unmarshal(in.Body, &doc); err != nil { + return nil, false + } + injected := false + if wantTags { + var meta map[string]any + if existing, ok := doc["metadata"]; ok { + if typed, isObject := existing.(map[string]any); isObject { + meta = typed + } + // non-object metadata: leave it; tags go unwritten so we + // don't clobber the client's value. Header fallback covers + // spend tracking. + } else { + meta = map[string]any{} + } + if meta != nil { + meta["tags"] = tags + doc["metadata"] = meta + injected = true + } + } + if wantUser { + // Anti-spoof: overwrite any client-supplied "user" so the + // gateway only sees our trusted identity. + doc["user"] = userID + injected = true + } + if !injected { + return nil, false + } + out, err := json.Marshal(doc) + if err != nil { + return nil, false + } + return out, true +} + +// applyJSONMetadata builds the Portkey-style mutations: a single header +// carrying a JSON object keyed by the rule's reserved field names. Per- +// value byte length is capped at MaxValueLength when set (Portkey +// enforces 128 chars). +func applyJSONMetadata(rule *JSONMetadataRule, in *middleware.Input) *middleware.Mutations { + mutations := &middleware.Mutations{} + mutations.HeadersRemove = append(mutations.HeadersRemove, rule.Header) + + payload := map[string]string{} + if rule.UserKey != "" { + if identity := identityFor(in); identity != "" { + payload[rule.UserKey] = truncate(identity, rule.MaxValueLength) + } + } + if rule.GroupsKey != "" { + if csv := authorisingTagsCSV(in); csv != "" { + payload[rule.GroupsKey] = truncate(csv, rule.MaxValueLength) + } + } + if len(payload) == 0 { + return mutations + } + raw, err := json.Marshal(payload) + if err != nil { + return mutations + } + mutations.HeadersAdd = append(mutations.HeadersAdd, middleware.KV{ + Key: rule.Header, + Value: string(raw), + }) + return mutations +} + +// identityFor returns the caller's display identity. UserEmail wins +// (carries the user email when peer-attached, peer.Name otherwise); +// UserID falls in only as a defensive last resort. +func identityFor(in *middleware.Input) string { + if in.UserEmail != "" { + return in.UserEmail + } + return in.UserID +} + +// authorisingTagsSlice returns the sorted, deduplicated slice of group +// display names the request was authorised under. Prefers the per- +// request authorising groups emitted by llm_router (intersection of the +// caller's UserGroups with the resolved route's AllowedGroupIDs) so the +// tags carry only the groups that actually authorise THIS request, not +// every group the peer happens to be in. Falls back to the full +// UserGroups when the router metadata key is absent. +func authorisingTagsSlice(in *middleware.Input) []string { + ids := tagsIDsFromAuthorising(in.Metadata) + if len(ids) == 0 { + ids = in.UserGroups + } + return tagsNamedSlice(ids, in.UserGroups, in.UserGroupNames) +} + +// authorisingTagsCSV is a convenience wrapper that joins +// authorisingTagsSlice with commas for HeaderPair-style emission. +func authorisingTagsCSV(in *middleware.Input) string { + return strings.Join(authorisingTagsSlice(in), ",") +} + +// truncate caps s to maxBytes bytes when maxBytes > 0. No-op when +// maxBytes <= 0 or s already fits. Truncation is byte-wise — sufficient +// for Portkey's 128-char ASCII limit. UTF-8 sequences could in theory +// be split, but the gateway treats the value as opaque bytes. +func truncate(s string, maxBytes int) string { + if maxBytes <= 0 || len(s) <= maxBytes { + return s + } + return s[:maxBytes] +} + +// tagsIDsFromAuthorising reads llm_router's authorising-groups metadata +// (a CSV of group ids) and returns the parsed slice. Returns nil when +// the key is absent or empty so the caller can fall back to the full +// UserGroups. +func tagsIDsFromAuthorising(meta []middleware.KV) []string { + v, ok := lookupMetadata(meta, middleware.KeyLLMAuthorisingGroups) + if !ok { + return nil + } + v = strings.TrimSpace(v) + if v == "" { + return nil + } + parts := strings.Split(v, ",") + out := make([]string, 0, len(parts)) + for _, p := range parts { + p = strings.TrimSpace(p) + if p != "" { + out = append(out, p) + } + } + if len(out) == 0 { + return nil + } + return out +} + +// tagsNamedSlice returns the sorted, deduplicated list of group display +// names. ids carries the canonical group identifiers to emit; +// userGroups + userGroupNames provide the positional id→name +// translation table from the Input envelope. When a name is missing +// for a given id (slice shorter than userGroups, or id absent from the +// table), the id is used verbatim so the tag still attributes +// correctly. Sorted so the same caller produces the same header value +// across requests (helps gateway-side cache hits and log correlation). +func tagsNamedSlice(ids, userGroups, userGroupNames []string) []string { + if len(ids) == 0 { + return nil + } + idToName := make(map[string]string, len(userGroups)) + for i, id := range userGroups { + if i < len(userGroupNames) { + idToName[id] = userGroupNames[i] + } + } + seen := make(map[string]struct{}, len(ids)) + out := make([]string, 0, len(ids)) + for _, id := range ids { + id = strings.TrimSpace(id) + if id == "" { + continue + } + tag := idToName[id] + if tag == "" { + tag = id + } + if _, dup := seen[tag]; dup { + continue + } + seen[tag] = struct{}{} + out = append(out, tag) + } + if len(out) == 0 { + return nil + } + sort.Strings(out) + return out +} + +// lookupMetadata returns the value for key plus a presence flag. +func lookupMetadata(meta []middleware.KV, key string) (string, bool) { + for _, kv := range meta { + if kv.Key == key { + return kv.Value, true + } + } + return "", false +} diff --git a/proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go b/proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go new file mode 100644 index 000000000..aab1271d8 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go @@ -0,0 +1,666 @@ +package llm_identity_inject + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +const ( + litellmProvider = "ainp_litellm-test" + portkeyProvider = "ainp_portkey-test" +) + +func newInput(resolvedProvider, userID string, groups []string) *middleware.Input { + return &middleware.Input{ + Slot: middleware.SlotOnRequest, + AccountID: "acct-test", + UserID: userID, + UserGroups: groups, + SourceIP: "100.64.0.5", + RequestID: "req-1", + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMResolvedProviderID, Value: resolvedProvider}, + }, + } +} + +func liteLLMRule() ProviderInjection { + return ProviderInjection{ + ProviderID: litellmProvider, + HeaderPair: &HeaderPairRule{ + EndUserIDHeader: "x-litellm-end-user-id", + TagsHeader: "x-litellm-tags", + }, + } +} + +func TestMiddlewareIdentity(t *testing.T) { + mw := New(Config{}) + assert.Equal(t, ID, mw.ID()) + assert.Equal(t, Version, mw.Version()) + assert.Equal(t, middleware.SlotOnRequest, mw.Slot()) + assert.True(t, mw.MutationsSupported()) + assert.Empty(t, mw.MetadataKeys(), "middleware emits no metadata") + assert.Nil(t, mw.AcceptedContentTypes()) + require.NoError(t, mw.Close()) +} + +func TestInject_MatchedProvider_StampsHeaders(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}}) + in := newInput(litellmProvider, "alice", []string{"grp-eng", "grp-it"}) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision) + require.NotNil(t, out.Mutations) + + // Strips the same headers we're about to add (anti-spoof). + assert.ElementsMatch(t, + []string{"x-litellm-end-user-id", "x-litellm-tags"}, + out.Mutations.HeadersRemove, + "every injected header must also appear in HeadersRemove so client-supplied values are wiped before our trusted values land") + + added := map[string]string{} + for _, kv := range out.Mutations.HeadersAdd { + added[kv.Key] = kv.Value + } + assert.Equal(t, "alice", added["x-litellm-end-user-id"]) + assert.Equal(t, "grp-eng,grp-it", added["x-litellm-tags"], "tags CSV must be sorted") +} + +func TestInject_UnmatchedProvider_NoMutations(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}}) + in := newInput("ainp_some-other-provider", "alice", []string{"grp-eng"}) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision) + assert.Nil(t, out.Mutations, "non-LiteLLM resolved provider must produce no mutations") +} + +func TestInject_NoResolvedProvider_NoMutations(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}}) + in := &middleware.Input{Slot: middleware.SlotOnRequest, UserID: "alice"} + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + assert.Nil(t, out.Mutations, + "missing llm.resolved_provider_id metadata means the router didn't run; never stamp identity blindly") +} + +func TestInject_PartialRule_StampsOnlyConfiguredHeaders(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{{ + ProviderID: litellmProvider, + HeaderPair: &HeaderPairRule{ + EndUserIDHeader: "x-litellm-end-user-id", + // TagsHeader intentionally empty. + }, + }}}) + in := newInput(litellmProvider, "alice", []string{"grp-eng"}) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + require.NotNil(t, out.Mutations) + + assert.Equal(t, []string{"x-litellm-end-user-id"}, out.Mutations.HeadersRemove, + "only configured header should be stripped") + require.Len(t, out.Mutations.HeadersAdd, 1) + assert.Equal(t, "x-litellm-end-user-id", out.Mutations.HeadersAdd[0].Key) + assert.Equal(t, "alice", out.Mutations.HeadersAdd[0].Value) +} + +func TestInject_EmptyIdentity_StripsButDoesNotAdd(t *testing.T) { + // Caller has no UserID and no groups. We still strip the headers + // (so the client can't inject identity) but we don't add empty + // values that would mislead the gateway. + mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}}) + in := newInput(litellmProvider, "", nil) + in.AccountID = "" + in.SourceIP = "" + in.RequestID = "" + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + require.NotNil(t, out.Mutations) + + assert.ElementsMatch(t, + []string{"x-litellm-end-user-id", "x-litellm-tags"}, + out.Mutations.HeadersRemove, + "identity headers must be stripped even when we don't have values to add — anti-spoof") + assert.Empty(t, out.Mutations.HeadersAdd, + "no NetBird identity available; do not stamp empty / misleading values") +} + +func TestInject_TagsCSV_DedupesAndSorts(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}}) + in := newInput(litellmProvider, "alice", []string{"grp-zzz", "grp-aaa", "grp-zzz", "", " "}) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + require.NotNil(t, out.Mutations) + + for _, kv := range out.Mutations.HeadersAdd { + if kv.Key == "x-litellm-tags" { + assert.Equal(t, "grp-aaa,grp-zzz", kv.Value, + "tags CSV must dedupe, drop empty, and sort") + return + } + } + t.Fatalf("expected x-litellm-tags in HeadersAdd; got %v", out.Mutations.HeadersAdd) +} + +func TestFactory_RejectsBadJSON(t *testing.T) { + _, err := Factory{}.New([]byte("{not json")) + require.Error(t, err) +} + +func TestFactory_AcceptsEmptyShapes(t *testing.T) { + for _, raw := range [][]byte{nil, []byte(""), []byte(" "), []byte("null"), []byte("{}"), []byte("[]")} { + mw, err := Factory{}.New(raw) + require.NoError(t, err) + require.NotNil(t, mw) + + out, ierr := mw.Invoke(context.Background(), + newInput(litellmProvider, "alice", []string{"grp-eng"})) + require.NoError(t, ierr) + assert.Equal(t, middleware.DecisionAllow, out.Decision) + assert.Nil(t, out.Mutations, + "empty config means no providers to inject for; every resolved provider passes through") + } +} + +func TestFactory_DropsInjectionRuleWithEmptyHeaders(t *testing.T) { + mw, err := Factory{}.New([]byte(`{"providers":[{"provider_id":"x"}]}`)) + require.NoError(t, err) + out, ierr := mw.Invoke(context.Background(), newInput("x", "alice", []string{"grp-eng"})) + require.NoError(t, ierr) + assert.Nil(t, out.Mutations, + "a rule with no header names is functionally a no-op and must be dropped at New() time") +} + +// TestInject_TagsFromAuthorisingMetadata pins that when llm_router has +// emitted llm.authorising_groups, the inject middleware uses THAT +// (the per-request authorising intersection) for the tags header — not +// the full UserGroups, which can include groups unrelated to this +// request's routing. +func TestInject_TagsFromAuthorisingMetadata(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}}) + in := newInput(litellmProvider, "alice", []string{"grp-eng", "grp-it", "grp-oncall"}) + in.Metadata = append(in.Metadata, middleware.KV{ + Key: middleware.KeyLLMAuthorisingGroups, + Value: "grp-eng", + }) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + require.NotNil(t, out.Mutations) + + for _, kv := range out.Mutations.HeadersAdd { + if kv.Key == "x-litellm-tags" { + assert.Equal(t, "grp-eng", kv.Value, + "tags must come from llm.authorising_groups, not the full UserGroups; unrelated peer groups must not leak") + return + } + } + t.Fatalf("expected x-litellm-tags in HeadersAdd; got %v", out.Mutations.HeadersAdd) +} + +// TestInject_TagsFallsBackToUserGroups pins the defensive fallback: if +// llm_router didn't emit authorising-groups metadata (chain +// misconfiguration) the middleware uses UserGroups so identity is +// still stamped, just over-broad. +func TestInject_TagsFallsBackToUserGroups(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}}) + in := newInput(litellmProvider, "alice", []string{"grp-eng", "grp-it"}) + // No llm.authorising_groups metadata. + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + require.NotNil(t, out.Mutations) + + for _, kv := range out.Mutations.HeadersAdd { + if kv.Key == "x-litellm-tags" { + assert.Equal(t, "grp-eng,grp-it", kv.Value, + "absent metadata must fall back to the full UserGroups CSV") + return + } + } + t.Fatalf("expected x-litellm-tags in HeadersAdd; got %v", out.Mutations.HeadersAdd) +} + +// portkeyRule is the JSONMetadata-shape analogue of liteLLMRule: a +// single x-portkey-metadata header carrying _user and groups, with +// Portkey's 128-byte per-value cap. +func portkeyRule() ProviderInjection { + return ProviderInjection{ + ProviderID: portkeyProvider, + JSONMetadata: &JSONMetadataRule{ + Header: "x-portkey-metadata", + UserKey: "_user", + GroupsKey: "groups", + MaxValueLength: 128, + }, + } +} + +// TestInject_JSONMetadata_StampsHeader pins the Portkey-style emission: +// one header carrying a JSON envelope with reserved keys for user +// identity and groups CSV. +func TestInject_JSONMetadata_StampsHeader(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{portkeyRule()}}) + in := newInput(portkeyProvider, "alice", []string{"grp-eng", "grp-it"}) + in.UserEmail = "alice@example.com" + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + require.NotNil(t, out.Mutations) + + assert.Equal(t, []string{"x-portkey-metadata"}, out.Mutations.HeadersRemove, + "the JSON header must be stripped before we add our trusted value") + require.Len(t, out.Mutations.HeadersAdd, 1) + added := out.Mutations.HeadersAdd[0] + assert.Equal(t, "x-portkey-metadata", added.Key) + + var payload map[string]string + require.NoError(t, json.Unmarshal([]byte(added.Value), &payload)) + assert.Equal(t, "alice@example.com", payload["_user"], + "_user reserved key carries the display identity (UserEmail)") + assert.Equal(t, "grp-eng,grp-it", payload["groups"], + "groups key carries the sorted CSV of group display names") +} + +// TestInject_JSONMetadata_TruncatesValues pins the per-value byte cap. +// Portkey rejects metadata values longer than 128 chars; oversized +// values are truncated rather than failing the request. +func TestInject_JSONMetadata_TruncatesValues(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{portkeyRule()}}) + in := newInput(portkeyProvider, "alice", []string{"grp-eng"}) + in.UserEmail = strings.Repeat("a", 200) + "@example.com" + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out.Mutations) + require.Len(t, out.Mutations.HeadersAdd, 1) + + var payload map[string]string + require.NoError(t, json.Unmarshal([]byte(out.Mutations.HeadersAdd[0].Value), &payload)) + assert.Len(t, payload["_user"], 128, + "per-value byte length must be capped at MaxValueLength") +} + +// TestInject_JSONMetadata_EmptyIdentity_StripsButDoesNotAdd verifies the +// anti-spoof Remove still fires when there's nothing to stamp. +func TestInject_JSONMetadata_EmptyIdentity_StripsButDoesNotAdd(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{portkeyRule()}}) + in := newInput(portkeyProvider, "", nil) + in.UserEmail = "" + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out.Mutations) + + assert.Equal(t, []string{"x-portkey-metadata"}, out.Mutations.HeadersRemove, + "strip even with no payload — client can't smuggle identity headers") + assert.Empty(t, out.Mutations.HeadersAdd, + "no NetBird identity available; do not stamp empty / misleading values") +} + +// TestFactory_RejectsRuleWithBothShapes pins the configuration-error +// guard: a rule that sets both HeaderPair and JSONMetadata is dropped +// at New() time rather than guessing which wins. +func TestFactory_RejectsRuleWithBothShapes(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{{ + ProviderID: litellmProvider, + HeaderPair: &HeaderPairRule{ + EndUserIDHeader: "x-litellm-end-user-id", + }, + JSONMetadata: &JSONMetadataRule{ + Header: "x-portkey-metadata", + UserKey: "_user", + }, + }}}) + in := newInput(litellmProvider, "alice", []string{"grp-eng"}) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + assert.Nil(t, out.Mutations, + "a rule that sets both shapes is ambiguous and must be dropped at New() time") +} + +// liteLLMRuleWithBody is the LiteLLM-style rule with body tag injection +// enabled (matches the catalog default). +func liteLLMRuleWithBody() ProviderInjection { + return ProviderInjection{ + ProviderID: litellmProvider, + HeaderPair: &HeaderPairRule{ + EndUserIDHeader: "x-litellm-end-user-id", + TagsHeader: "x-litellm-tags", + TagsInBody: true, + }, + } +} + +// TestInject_BodyTags_AddsMetadataTags pins the body-inject path that +// LiteLLM's _tag_max_budget_check requires. With TagsInBody set, the +// middleware writes the authorising-groups slice into +// request.metadata.tags (in addition to the header). +func TestInject_BodyTags_AddsMetadataTags(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}}) + in := newInput(litellmProvider, "alice", []string{"grp-eng", "grp-sre"}) + in.Body = []byte(`{"model":"gpt-4o-mini","messages":[]}`) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out.Mutations) + require.NotEmpty(t, out.Mutations.BodyReplace, "body must be rewritten when TagsInBody is set") + + var doc map[string]any + require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc)) + meta, ok := doc["metadata"].(map[string]any) + require.True(t, ok, "metadata must be an object") + tags, ok := meta["tags"].([]any) + require.True(t, ok, "metadata.tags must be a JSON array") + got := make([]string, 0, len(tags)) + for _, t := range tags { + s, _ := t.(string) + got = append(got, s) + } + assert.Equal(t, []string{"grp-eng", "grp-sre"}, got, + "metadata.tags must carry the sorted authorising-groups slice") + assert.Equal(t, "gpt-4o-mini", doc["model"], + "the rest of the body must be preserved verbatim") +} + +// TestInject_BodyTags_PreservesExistingMetadata pins that an existing +// metadata object on the request is merged with our tags rather than +// clobbered — clients sometimes set metadata fields the proxy +// shouldn't blow away (jobID, taskName, etc.). +func TestInject_BodyTags_PreservesExistingMetadata(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}}) + in := newInput(litellmProvider, "alice", []string{"grp-eng"}) + in.Body = []byte(`{"model":"gpt-4o-mini","metadata":{"jobID":"j-42","tags":["should-be-replaced"]}}`) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotEmpty(t, out.Mutations.BodyReplace) + + var doc map[string]any + require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc)) + meta := doc["metadata"].(map[string]any) + assert.Equal(t, "j-42", meta["jobID"], + "client-supplied metadata fields outside `tags` must survive") + tags := meta["tags"].([]any) + require.Len(t, tags, 1) + assert.Equal(t, "grp-eng", tags[0], + "our tags overwrite any client-supplied metadata.tags so spoofing is impossible") +} + +// TestInject_BodyTags_SkipsHostileMetadataShape pins the defensive +// refusal: when the request body has a non-object metadata field +// (string/number/array), we don't inject — header path still emits. +func TestInject_BodyTags_SkipsHostileMetadataShape(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}}) + in := newInput(litellmProvider, "alice", []string{"grp-eng"}) + in.Body = []byte(`{"model":"gpt-4o-mini","metadata":"not-an-object"}`) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out.Mutations) + assert.Empty(t, out.Mutations.BodyReplace, + "non-object metadata must skip body inject (don't clobber)") + + for _, kv := range out.Mutations.HeadersAdd { + if kv.Key == "x-litellm-tags" { + assert.Equal(t, "grp-eng", kv.Value, + "header path must still emit so spend tracking keeps working") + return + } + } + t.Fatalf("expected x-litellm-tags header even when body inject was skipped") +} + +// TestInject_BodyTags_SkipsTruncatedBody pins that we don't blindly +// rewrite a body we don't have in full. The header path still runs. +func TestInject_BodyTags_SkipsTruncatedBody(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}}) + in := newInput(litellmProvider, "alice", []string{"grp-eng"}) + in.Body = []byte(`{"model":"gpt-4o-mini","messages":[]}`) + in.BodyTruncated = true + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + assert.Empty(t, out.Mutations.BodyReplace, + "truncated body must skip body inject — re-marshaling would corrupt the request") +} + +// TestInject_BodyTags_SkipsNonJSONBody pins graceful behavior when the +// body isn't JSON (e.g. a streaming binary or form upload sneaking +// through the LLM chain). Header path still runs. +func TestInject_BodyTags_SkipsNonJSONBody(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}}) + in := newInput(litellmProvider, "alice", []string{"grp-eng"}) + in.Body = []byte(`not even close to json`) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + assert.Empty(t, out.Mutations.BodyReplace, + "non-JSON body must skip body inject silently") +} + +// liteLLMRuleFull mirrors the catalog default: header path + body +// metadata.tags (groups) + body user (end-user id). +func liteLLMRuleFull() ProviderInjection { + return ProviderInjection{ + ProviderID: litellmProvider, + HeaderPair: &HeaderPairRule{ + EndUserIDHeader: "x-litellm-end-user-id", + TagsHeader: "x-litellm-tags", + TagsInBody: true, + EndUserIDInBody: true, + }, + } +} + +// TestInject_BodyUser_WritesTopLevelUser pins the EndUserIDInBody path +// alone: body's top-level "user" field carries the display identity. +// Tags-in-body is OFF here so we isolate the user write. +func TestInject_BodyUser_WritesTopLevelUser(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{{ + ProviderID: litellmProvider, + HeaderPair: &HeaderPairRule{ + EndUserIDHeader: "x-litellm-end-user-id", + EndUserIDInBody: true, + }, + }}}) + in := newInput(litellmProvider, "alice", nil) + in.UserEmail = "alice@example.com" + in.Body = []byte(`{"model":"gpt-4o-mini","messages":[]}`) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out.Mutations) + require.NotEmpty(t, out.Mutations.BodyReplace) + + var doc map[string]any + require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc)) + assert.Equal(t, "alice@example.com", doc["user"], + "body's top-level user field must carry the display identity") + _, hasMeta := doc["metadata"] + assert.False(t, hasMeta, "TagsInBody is off; metadata must not be added") +} + +// TestInject_BodyUser_OverwritesClientSupplied pins anti-spoof: a +// client-supplied "user" in the body is overwritten so the gateway +// only sees our trusted identity. +func TestInject_BodyUser_OverwritesClientSupplied(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRuleFull()}}) + in := newInput(litellmProvider, "alice", []string{"grp-eng"}) + in.UserEmail = "alice@example.com" + in.Body = []byte(`{"model":"gpt-4o-mini","user":"ceo@company.com"}`) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotEmpty(t, out.Mutations.BodyReplace) + + var doc map[string]any + require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc)) + assert.Equal(t, "alice@example.com", doc["user"], + "client-supplied user must be overwritten with the trusted identity") +} + +// TestInject_BodyCombined_TagsAndUser pins that with both flags on, +// the body carries both metadata.tags AND top-level user, and the +// header path still emits. +func TestInject_BodyCombined_TagsAndUser(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRuleFull()}}) + in := newInput(litellmProvider, "alice", []string{"grp-eng", "grp-sre"}) + in.UserEmail = "alice@example.com" + in.Body = []byte(`{"model":"gpt-4o-mini"}`) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotEmpty(t, out.Mutations.BodyReplace) + + var doc map[string]any + require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc)) + assert.Equal(t, "alice@example.com", doc["user"]) + meta := doc["metadata"].(map[string]any) + tags := meta["tags"].([]any) + require.Len(t, tags, 2) + assert.Equal(t, "grp-eng", tags[0]) + assert.Equal(t, "grp-sre", tags[1]) + + // Header path still emits — header end-user-id is the primary + // path for LiteLLM's resolver, body is defense-in-depth. + added := map[string]string{} + for _, kv := range out.Mutations.HeadersAdd { + added[kv.Key] = kv.Value + } + assert.Equal(t, "alice@example.com", added["x-litellm-end-user-id"]) + assert.Equal(t, "grp-eng,grp-sre", added["x-litellm-tags"]) +} + +// TestInject_BodyCombined_HostileMetadataKeepsUser pins the partial- +// success path: a hostile (non-object) metadata field skips the tag +// write but still allows the orthogonal user write to land. +func TestInject_BodyCombined_HostileMetadataKeepsUser(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRuleFull()}}) + in := newInput(litellmProvider, "alice", []string{"grp-eng"}) + in.UserEmail = "alice@example.com" + in.Body = []byte(`{"model":"gpt-4o-mini","metadata":"not-an-object"}`) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotEmpty(t, out.Mutations.BodyReplace, + "user write must still go through even when metadata is hostile") + + var doc map[string]any + require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc)) + assert.Equal(t, "alice@example.com", doc["user"]) + assert.Equal(t, "not-an-object", doc["metadata"], + "hostile metadata must be left untouched, not clobbered") +} + +// TestInject_ExtraHeaders_Stamped pins the extras path: with a +// per-provider ExtraHeader configured (e.g. Portkey config id), the +// middleware stamps it on every matching request and adds the same +// name to HeadersRemove for anti-spoof. +func TestInject_ExtraHeaders_Stamped(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{{ + ProviderID: portkeyProvider, + JSONMetadata: &JSONMetadataRule{ + Header: "x-portkey-metadata", + UserKey: "_user", + GroupsKey: "groups", + }, + ExtraHeaders: []ExtraHeaderKV{ + {Name: "x-portkey-config", Value: "pc-prod-3f2a"}, + }, + }}}) + in := newInput(portkeyProvider, "alice", []string{"grp-eng"}) + in.UserEmail = "alice@example.com" + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out.Mutations) + + assert.Contains(t, out.Mutations.HeadersRemove, "x-portkey-config", + "extras must be stripped before stamping for anti-spoof") + added := map[string]string{} + for _, kv := range out.Mutations.HeadersAdd { + added[kv.Key] = kv.Value + } + assert.Equal(t, "pc-prod-3f2a", added["x-portkey-config"], + "extras must carry the operator-configured value verbatim") + // Identity-stamping shape (JSONMetadata header) still emitted. + assert.Contains(t, added, "x-portkey-metadata", + "extras and identity stamping are independent — both must land") +} + +// TestInject_ExtraHeaders_OnlyRule pins that an extras-only rule +// (no HeaderPair, no JSONMetadata) survives New() and stamps the +// extras anyway. Useful for hypothetical gateways that need a static +// routing header but no NetBird identity stamping. +func TestInject_ExtraHeaders_OnlyRule(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{{ + ProviderID: "ainp_extras-only", + ExtraHeaders: []ExtraHeaderKV{ + {Name: "x-routing-key", Value: "rk-1"}, + }, + }}}) + in := newInput("ainp_extras-only", "alice", nil) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out.Mutations, + "extras alone keep the rule alive — middleware must emit them") + added := map[string]string{} + for _, kv := range out.Mutations.HeadersAdd { + added[kv.Key] = kv.Value + } + assert.Equal(t, "rk-1", added["x-routing-key"]) +} + +// TestInject_ExtraHeaders_EmptyValueSkipped pins that empty values are +// dropped silently (the synth would normally not send them, but the +// middleware is defensive). +func TestInject_ExtraHeaders_EmptyValueSkipped(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{{ + ProviderID: portkeyProvider, + JSONMetadata: &JSONMetadataRule{ + Header: "x-portkey-metadata", + UserKey: "_user", + }, + ExtraHeaders: []ExtraHeaderKV{ + {Name: "x-portkey-config", Value: ""}, + }, + }}}) + in := newInput(portkeyProvider, "alice", []string{"grp-eng"}) + in.UserEmail = "alice@example.com" + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out.Mutations) + assert.NotContains(t, out.Mutations.HeadersRemove, "x-portkey-config", + "empty extra value must not even strip the header") + for _, kv := range out.Mutations.HeadersAdd { + assert.NotEqual(t, "x-portkey-config", kv.Key, + "empty extra value must not be stamped") + } +} diff --git a/proxy/internal/middleware/builtin/llm_limit_check/factory.go b/proxy/internal/middleware/builtin/llm_limit_check/factory.go new file mode 100644 index 000000000..1068a6867 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_limit_check/factory.go @@ -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{}) +} diff --git a/proxy/internal/middleware/builtin/llm_limit_check/middleware.go b/proxy/internal/middleware/builtin/llm_limit_check/middleware.go new file mode 100644 index 000000000..bebe4dca4 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_limit_check/middleware.go @@ -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 "" +} diff --git a/proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go b/proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go new file mode 100644 index 000000000..2c26c2abe --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go @@ -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) +} diff --git a/proxy/internal/middleware/builtin/llm_limit_record/factory.go b/proxy/internal/middleware/builtin/llm_limit_record/factory.go new file mode 100644 index 000000000..b42931c0c --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_limit_record/factory.go @@ -0,0 +1,35 @@ +// Package llm_limit_record is the SlotOnResponse middleware that +// posts the served request's token + cost deltas back to management +// so the per-(user, group, window) consumption counters tick. Reads +// the attribution metadata stamped by llm_limit_check on the request +// leg + the token / cost metadata stamped by llm_response_parser and +// cost_meter; skips the write entirely when no attribution metadata +// is present (e.g. catch-all-allow policy with no caps configured). +package llm_limit_record + +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_record" + +// Factory builds a configured llm_limit_record instance bound to the +// FactoryContext's MgmtClient. nil-MgmtClient disables the post-flight +// write entirely (no-op pass-through), matching the request-leg gate's +// behaviour so a partially wired environment is consistent. +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). +func (Factory) New(_ []byte) (middleware.Middleware, error) { + ctx := builtin.Context() + return New(ctx.MgmtClient, ctx.Logger), nil +} + +func init() { + builtin.Register(Factory{}) +} diff --git a/proxy/internal/middleware/builtin/llm_limit_record/middleware.go b/proxy/internal/middleware/builtin/llm_limit_record/middleware.go new file mode 100644 index 000000000..52dfc73f0 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_limit_record/middleware.go @@ -0,0 +1,144 @@ +package llm_limit_record + +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 post-flight RPC. +// Longer than the pre-flight gate because this runs after the +// upstream returned and is not on the user-facing latency path — +// a slow record is just a delayed counter increment, not a delayed +// response to the caller. +const callTimeout = 5 * time.Second + +// Middleware posts token + cost deltas to management after a served +// request. Stateless; per-call values come entirely from metadata +// emitted upstream. +type Middleware struct { + mgmt builtin.MgmtClient + logger *log.Logger +} + +// New constructs a Middleware bound to the supplied management +// client. mgmt may be nil — that disables the write entirely so a +// partially wired environment doesn't attempt to dial nothing. +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 that the middleware runs after the upstream call. +func (m *Middleware) Slot() middleware.Slot { return middleware.SlotOnResponse } + +// AcceptedContentTypes is empty: this middleware never inspects +// bodies. It only reads metadata emitted upstream. +func (m *Middleware) AcceptedContentTypes() []string { return []string{} } + +// MetadataKeys is empty — the record middleware never emits its own +// metadata. Its only side effect is the gRPC write to management. +func (m *Middleware) MetadataKeys() []string { return []string{} } + +// MutationsSupported reports that the middleware never mutates the +// response. Its outcome is always Allow. +func (m *Middleware) MutationsSupported() bool { return false } + +// Close releases resources owned by the middleware. Stateless. +func (m *Middleware) Close() error { return nil } + +// Invoke reads the attribution + tokens + cost metadata, calls +// management's RecordLLMUsage, and always returns Allow. RPC errors +// are logged at debug level — the response has already been served +// to the client by the time we get here, so a record failure must +// not surface back through the proxy. +func (m *Middleware) Invoke(ctx context.Context, in *middleware.Input) (*middleware.Output, error) { + out := &middleware.Output{Decision: middleware.DecisionAllow} + if m.mgmt == nil { + return out, nil + } + + tokensIn, _ := strconv.ParseInt(lookupKV(in.Metadata, middleware.KeyLLMInputTokens), 10, 64) + tokensOut, _ := strconv.ParseInt(lookupKV(in.Metadata, middleware.KeyLLMOutputTokens), 10, 64) + costUSD, _ := strconv.ParseFloat(lookupKV(in.Metadata, middleware.KeyCostUSDTotal), 64) + if tokensIn == 0 && tokensOut == 0 && costUSD == 0 { + // llm_response_parser couldn't read usage off the upstream + // response (streaming-not-yet-supported, malformed body, …). + // Skipping the write keeps phantom rows out of the + // consumption table. + return out, nil + } + + windowStr := lookupKV(in.Metadata, middleware.KeyLLMAttributionWindowS) + windowSeconds, _ := strconv.ParseInt(windowStr, 10, 64) + groupID := lookupKV(in.Metadata, middleware.KeyLLMAttributionGroupID) + + // A zero attribution window means no policy cap bound this request (deny at + // the gate, or a catch-all-allow policy). We still record so account-level + // budget rules — which live in their own windows and bind independently of + // policies — accumulate. The management side books the policy dimensions + // only when window_seconds > 0 and fans out to account rules regardless. + if in.UserID == "" && groupID == "" && len(in.UserGroups) == 0 { + m.logger.WithField("middleware", ID). + WithField("account_id", in.AccountID). + Debugf("post-flight skipped: no user/group/groups to attribute (tokens=%d/%d cost=%g window=%d)", tokensIn, tokensOut, costUSD, windowSeconds) + return out, nil + } + + rpcCtx, cancel := context.WithTimeout(ctx, callTimeout) + defer cancel() + + m.logger.WithField("middleware", ID). + WithField("account_id", in.AccountID). + WithField("user_id", in.UserID). + WithField("group_id", groupID). + WithField("group_ids_len", len(in.UserGroups)). + Debugf("post-flight sending RecordLLMUsage (tokens=%d/%d cost=%g window=%d)", tokensIn, tokensOut, costUSD, windowSeconds) + + if _, err := m.mgmt.RecordLLMUsage(rpcCtx, &proto.RecordLLMUsageRequest{ + AccountId: in.AccountID, + UserId: in.UserID, + GroupId: groupID, + WindowSeconds: windowSeconds, + TokensInput: tokensIn, + TokensOutput: tokensOut, + CostUsd: costUSD, + GroupIds: append([]string(nil), in.UserGroups...), + }); err != nil { + m.logger.WithError(err). + WithField("middleware", ID). + WithField("account_id", in.AccountID). + WithField("user_id", in.UserID). + WithField("group_id", groupID). + Debugf("post-flight record failed; counter will lag this request") + } + return out, nil +} + +// 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 "" +} diff --git a/proxy/internal/middleware/builtin/llm_limit_record/middleware_test.go b/proxy/internal/middleware/builtin/llm_limit_record/middleware_test.go new file mode 100644 index 000000000..a98ce9b9f --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_limit_record/middleware_test.go @@ -0,0 +1,191 @@ +package llm_limit_record + +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" +) + +type fakeMgmt struct { + recordReq *proto.RecordLLMUsageRequest + recordCalled bool + recordErr error +} + +func (f *fakeMgmt) CheckLLMPolicyLimits(_ context.Context, _ *proto.CheckLLMPolicyLimitsRequest, _ ...grpc.CallOption) (*proto.CheckLLMPolicyLimitsResponse, error) { + return &proto.CheckLLMPolicyLimitsResponse{Decision: "allow"}, nil +} + +func (f *fakeMgmt) RecordLLMUsage(_ context.Context, in *proto.RecordLLMUsageRequest, _ ...grpc.CallOption) (*proto.RecordLLMUsageResponse, error) { + f.recordCalled = true + f.recordReq = in + return &proto.RecordLLMUsageResponse{}, f.recordErr +} + +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) + require.NotNil(t, out) + return out +} + +// TestInvoke_PostsAttributionWithTokensAndCost covers the happy path: +// when the request leg stamped attribution + the upstream parsers +// stamped tokens + cost, the post-flight call carries every field +// through to RecordLLMUsage. +func TestInvoke_PostsAttributionWithTokensAndCost(t *testing.T) { + mgmt := &fakeMgmt{} + m := New(mgmt, nil) + + out := runInvoke(t, m, &middleware.Input{ + AccountID: "acc-1", + UserID: "user-bob", + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMAttributionGroupID, Value: "grp-engineers"}, + {Key: middleware.KeyLLMAttributionWindowS, Value: "86400"}, + {Key: middleware.KeyLLMInputTokens, Value: "150"}, + {Key: middleware.KeyLLMOutputTokens, Value: "75"}, + {Key: middleware.KeyCostUSDTotal, Value: "0.0125"}, + }, + }) + + assert.Equal(t, middleware.DecisionAllow, out.Decision) + require.True(t, mgmt.recordCalled, "record must be invoked when attribution + usage are both present") + assert.Equal(t, "acc-1", mgmt.recordReq.GetAccountId()) + assert.Equal(t, "user-bob", mgmt.recordReq.GetUserId()) + assert.Equal(t, "grp-engineers", mgmt.recordReq.GetGroupId()) + assert.Equal(t, int64(86_400), mgmt.recordReq.GetWindowSeconds()) + assert.Equal(t, int64(150), mgmt.recordReq.GetTokensInput()) + assert.Equal(t, int64(75), mgmt.recordReq.GetTokensOutput()) + assert.InDelta(t, 0.0125, mgmt.recordReq.GetCostUsd(), 1e-9) +} + +// TestInvoke_NoAttributionWindowStillRecordsForAccountFanOut proves the +// catch-all-allow path now STILL records (window 0): account-level budget +// rules live in their own windows and bind independently of policies, so the +// management side needs the post-flight call even when no policy cap applied. +// The full group set is forwarded so the account fan-out can attribute. +func TestInvoke_NoAttributionWindowStillRecordsForAccountFanOut(t *testing.T) { + mgmt := &fakeMgmt{} + m := New(mgmt, nil) + + runInvoke(t, m, &middleware.Input{ + AccountID: "acc-1", + UserID: "user-bob", + UserGroups: []string{"grp-eng", "grp-oncall"}, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMInputTokens, Value: "150"}, + {Key: middleware.KeyLLMOutputTokens, Value: "75"}, + {Key: middleware.KeyCostUSDTotal, Value: "0.0125"}, + }, + }) + + require.True(t, mgmt.recordCalled, "must record even without a policy window so account budgets accumulate") + assert.Equal(t, int64(0), mgmt.recordReq.GetWindowSeconds(), "no policy window is forwarded as 0") + assert.Empty(t, mgmt.recordReq.GetGroupId(), "no attribution group without a policy") + assert.Equal(t, []string{"grp-eng", "grp-oncall"}, mgmt.recordReq.GetGroupIds(), "full group set must be forwarded for the account fan-out") +} + +// TestInvoke_NoPrincipalSkipsRecord proves that with neither a user nor any +// groups there is nothing to attribute, so the write is skipped. +func TestInvoke_NoPrincipalSkipsRecord(t *testing.T) { + mgmt := &fakeMgmt{} + m := New(mgmt, nil) + + runInvoke(t, m, &middleware.Input{ + AccountID: "acc-1", + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMInputTokens, Value: "150"}, + {Key: middleware.KeyCostUSDTotal, Value: "0.0125"}, + }, + }) + + assert.False(t, mgmt.recordCalled, "no user and no groups = nothing to attribute") +} + +// TestInvoke_ZeroUsageSkipsRecord proves the no-usage-no-write path: +// when the upstream parser couldn't extract token counts (streaming, +// malformed body, …), skipping the write keeps phantom rows out of +// the consumption table. +func TestInvoke_ZeroUsageSkipsRecord(t *testing.T) { + mgmt := &fakeMgmt{} + m := New(mgmt, nil) + + runInvoke(t, m, &middleware.Input{ + AccountID: "acc-1", + UserID: "user-bob", + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMAttributionGroupID, Value: "grp-engineers"}, + {Key: middleware.KeyLLMAttributionWindowS, Value: "86400"}, + }, + }) + + assert.False(t, mgmt.recordCalled, "zero tokens AND zero cost = nothing to record; an upstream parse miss must not surface as a row") +} + +// TestInvoke_RPCErrorIsSwallowed proves the post-flight isolation +// contract: management errors must NOT cascade back to the proxy +// because the upstream response has already been served — failing +// the chain at this point would corrupt the response. Errors are +// logged at debug level and swallowed. +func TestInvoke_RPCErrorIsSwallowed(t *testing.T) { + mgmt := &fakeMgmt{recordErr: errors.New("management down")} + m := New(mgmt, nil) + + out := runInvoke(t, m, &middleware.Input{ + AccountID: "acc-1", + UserID: "user-bob", + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMAttributionGroupID, Value: "grp-engineers"}, + {Key: middleware.KeyLLMAttributionWindowS, Value: "86400"}, + {Key: middleware.KeyLLMInputTokens, Value: "100"}, + }, + }) + + assert.Equal(t, middleware.DecisionAllow, out.Decision, + "a record failure must not surface — the upstream response is already on the wire") +} + +// TestInvoke_NoMgmtClientPassesThrough mirrors the gate's safety +// contract: a partial wiring is consistent. No mgmt client = silent +// skip rather than an unhandled nil-deref. +func TestInvoke_NoMgmtClientPassesThrough(t *testing.T) { + m := New(nil, nil) + out := runInvoke(t, m, &middleware.Input{ + AccountID: "acc-1", + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMAttributionGroupID, Value: "grp-engineers"}, + {Key: middleware.KeyLLMAttributionWindowS, Value: "86400"}, + {Key: middleware.KeyLLMInputTokens, Value: "100"}, + }, + }) + assert.Equal(t, middleware.DecisionAllow, out.Decision) +} + +// TestInvoke_NoIdentitySkipsRecord covers a defensive guard: stamped +// attribution but no user_id AND no group_id (shouldn't happen, but +// possible if the gate ever changes shape) must not write a row keyed +// on empty dimension ids. +func TestInvoke_NoIdentitySkipsRecord(t *testing.T) { + mgmt := &fakeMgmt{} + m := New(mgmt, nil) + + runInvoke(t, m, &middleware.Input{ + AccountID: "acc-1", + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMAttributionWindowS, Value: "86400"}, + {Key: middleware.KeyLLMInputTokens, Value: "100"}, + }, + }) + + assert.False(t, mgmt.recordCalled, + "empty user + group identity must skip the write — never key on empty dimension ids") +} diff --git a/proxy/internal/middleware/builtin/llm_request_parser/bedrock_test.go b/proxy/internal/middleware/builtin/llm_request_parser/bedrock_test.go new file mode 100644 index 000000000..827b81d07 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_request_parser/bedrock_test.go @@ -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) + } + } +} diff --git a/proxy/internal/middleware/builtin/llm_request_parser/factory.go b/proxy/internal/middleware/builtin/llm_request_parser/factory.go new file mode 100644 index 000000000..8b3776877 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_request_parser/factory.go @@ -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{}) +} diff --git a/proxy/internal/middleware/builtin/llm_request_parser/middleware.go b/proxy/internal/middleware/builtin/llm_request_parser/middleware.go new file mode 100644 index 000000000..64ca04e6a --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_request_parser/middleware.go @@ -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 +} diff --git a/proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go b/proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go new file mode 100644 index 000000000..bc185b295 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go @@ -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") +} diff --git a/proxy/internal/middleware/builtin/llm_response_parser/factory.go b/proxy/internal/middleware/builtin/llm_response_parser/factory.go new file mode 100644 index 000000000..e7d634109 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_response_parser/factory.go @@ -0,0 +1,43 @@ +package llm_response_parser + +import ( + "bytes" + "encoding/json" + "fmt" + + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/proxy/internal/middleware/builtin" +) + +// Factory constructs configured Middleware instances for the registry. +type Factory struct{} + +// ID returns the registry identifier. +func (Factory) ID() string { return ID } + +// New decodes RawConfig (empty / null / "{}" all accepted) and returns +// a configured Middleware. Construction never fails on a well-formed +// empty config; only structurally invalid JSON is rejected. +func (Factory) New(rawConfig []byte) (middleware.Middleware, error) { + cfg, err := decodeConfig(rawConfig) + if err != nil { + return nil, fmt.Errorf("decode config: %w", err) + } + return New(cfg), nil +} + +func decodeConfig(raw []byte) (config, error) { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { + return config{}, nil + } + var cfg config + if err := json.Unmarshal(trimmed, &cfg); err != nil { + return config{}, err + } + return cfg, nil +} + +func init() { + builtin.Register(Factory{}) +} diff --git a/proxy/internal/middleware/builtin/llm_response_parser/gzip_test.go b/proxy/internal/middleware/builtin/llm_response_parser/gzip_test.go new file mode 100644 index 000000000..017153e4c --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_response_parser/gzip_test.go @@ -0,0 +1,133 @@ +package llm_response_parser + +import ( + "bytes" + "compress/flate" + "compress/gzip" + "compress/zlib" + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +// gzipBytes returns data gzip-compressed — the wire shape Anthropic +// returns when the client (Claude Code) negotiated Accept-Encoding: gzip. +func gzipBytes(t *testing.T, data []byte) []byte { + t.Helper() + var buf bytes.Buffer + w := gzip.NewWriter(&buf) + _, err := w.Write(data) + require.NoError(t, err, "gzip write must succeed") + require.NoError(t, w.Close(), "gzip close must succeed") + return buf.Bytes() +} + +// TestInvoke_AnthropicStreaming_Gzip is the regression guard for the live +// bug: Claude Code negotiates gzip, Anthropic gzips the SSE stream, the +// proxy captures the compressed bytes, and the parser must decompress +// before accumulating — otherwise token usage is silently dropped and +// cost_meter skips with missing_tokens. +func TestInvoke_AnthropicStreaming_Gzip(t *testing.T) { + m := newTestMiddleware(t) + body := gzipBytes(t, loadFixture(t, "anthropic_stream.txt")) + + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{ + {Key: "Content-Type", Value: "text/event-stream; charset=utf-8"}, + {Key: "Content-Encoding", Value: "gzip"}, + }, + RespBody: body, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "anthropic"}, + {Key: middleware.KeyLLMModel, Value: "claude-opus-4-8"}, + }, + } + + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "Invoke must not error on a gzip-encoded streaming body") + + in123, ok := metaValue(out.Metadata, middleware.KeyLLMInputTokens) + require.True(t, ok, "input tokens must be emitted from a gzip SSE stream") + assert.Equal(t, "123", in123, "input tokens must survive gzip decompression") + + outTok, _ := metaValue(out.Metadata, middleware.KeyLLMOutputTokens) + assert.Equal(t, "45", outTok, "output tokens must survive gzip decompression") + + totTok, _ := metaValue(out.Metadata, middleware.KeyLLMTotalTokens) + assert.Equal(t, "168", totTok, "total tokens must survive gzip decompression") +} + +// TestInvoke_AnthropicBuffered_Gzip covers the non-streaming JSON path +// under gzip — the same decode must happen before ParseResponse. +func TestInvoke_AnthropicBuffered_Gzip(t *testing.T) { + m := newTestMiddleware(t) + body := gzipBytes(t, loadFixture(t, "anthropic_messages.json")) + + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{ + {Key: "Content-Type", Value: "application/json"}, + {Key: "Content-Encoding", Value: "gzip"}, + }, + RespBody: body, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "anthropic"}, + {Key: middleware.KeyLLMModel, Value: "claude-opus-4-8"}, + }, + } + + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "Invoke must not error on a gzip-encoded buffered body") + + _, ok := metaValue(out.Metadata, middleware.KeyLLMInputTokens) + require.True(t, ok, "input tokens must be emitted from a gzip JSON body") +} + +// TestDecodeResponseBody covers the encoding matrix directly. +func TestDecodeResponseBody(t *testing.T) { + plain := []byte(`{"hello":"world"}`) + + t.Run("identity passthrough", func(t *testing.T) { + assert.Equal(t, plain, decodeResponseBody(plain, "")) + assert.Equal(t, plain, decodeResponseBody(plain, "identity")) + }) + + t.Run("gzip", func(t *testing.T) { + assert.Equal(t, plain, decodeResponseBody(gzipBytes(t, plain), "gzip")) + }) + + t.Run("gzip with multi-coding header takes outermost", func(t *testing.T) { + assert.Equal(t, plain, decodeResponseBody(gzipBytes(t, plain), "identity, gzip")) + }) + + t.Run("deflate zlib-wrapped", func(t *testing.T) { + var buf bytes.Buffer + zw := zlib.NewWriter(&buf) + _, _ = zw.Write(plain) + _ = zw.Close() + assert.Equal(t, plain, decodeResponseBody(buf.Bytes(), "deflate")) + }) + + t.Run("deflate raw flate fallback", func(t *testing.T) { + var buf bytes.Buffer + fw, _ := flate.NewWriter(&buf, flate.DefaultCompression) + _, _ = fw.Write(plain) + _ = fw.Close() + assert.Equal(t, plain, decodeResponseBody(buf.Bytes(), "deflate")) + }) + + t.Run("gzip header but not actually gzip falls back to raw", func(t *testing.T) { + assert.Equal(t, plain, decodeResponseBody(plain, "gzip")) + }) + + t.Run("unknown encoding (br) returns raw", func(t *testing.T) { + assert.Equal(t, plain, decodeResponseBody(plain, "br")) + }) +} diff --git a/proxy/internal/middleware/builtin/llm_response_parser/middleware.go b/proxy/internal/middleware/builtin/llm_response_parser/middleware.go new file mode 100644 index 000000000..a204460cb --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_response_parser/middleware.go @@ -0,0 +1,339 @@ +// Package llm_response_parser implements the SlotOnResponse middleware +// that decodes OpenAI- and Anthropic-shaped LLM responses (buffered or +// streaming) and emits token usage and completion metadata. Provider +// and model are read from the request-side metadata bag emitted by +// llm_request_parser; without that context the middleware is a no-op. +package llm_response_parser + +import ( + "bytes" + "compress/flate" + "compress/gzip" + "compress/zlib" + "context" + "io" + "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/llm_guardrail" +) + +// ID is the registry identifier for this middleware. +const ID = "llm_response_parser" + +const version = "1.0.0" + +// maxCompletionBytes is the rune-safe cap applied to the extracted +// completion text before emitting it as metadata. +const maxCompletionBytes = 3500 + +// maxDecodedBytes bounds the inflated size of a compressed response body +// so a small gzip/deflate payload can't expand into a memory blow-up. The +// captured input is already capped (per-direction body cap), so this only +// bounds the decompression ratio; the parser is best-effort and tolerates a +// truncated decode. +const maxDecodedBytes = 16 << 20 + +var ( + acceptedContentTypes = []string{"application/json", "text/event-stream"} + metadataKeys = []string{ + middleware.KeyLLMInputTokens, + middleware.KeyLLMOutputTokens, + middleware.KeyLLMTotalTokens, + middleware.KeyLLMCachedInputTokens, + middleware.KeyLLMCacheCreationTokens, + middleware.KeyLLMResponseCompletion, + } +) + +// config is the wire-side configuration for this middleware. RedactPii, when +// true, runs PII redaction on the extracted completion text BEFORE it is +// emitted as llm.response_completion — keeping the access-log row free of +// emails / SSNs / phone numbers the model itself generated. CaptureCompletion +// gates emission of the completion key entirely: a nil pointer preserves +// legacy emission (so callers without the toggle aren't broken), an explicit +// false suppresses the key so the access-log row carries token / cost facts +// only. Both are sourced by the synthesiser from the account's redact_pii +// and enable_prompt_collection toggles respectively. +type config struct { + RedactPii bool `json:"redact_pii,omitempty"` + CaptureCompletion *bool `json:"capture_completion,omitempty"` +} + +// Middleware implements middleware.Middleware. +type Middleware struct { + parsers []llm.Parser + redactPii bool + captureCompletion bool +} + +// New constructs a configured Middleware instance. +func New(cfg config) *Middleware { + capture := true + if cfg.CaptureCompletion != nil { + capture = *cfg.CaptureCompletion + } + return &Middleware{parsers: llm.Parsers(), redactPii: cfg.RedactPii, captureCompletion: capture} +} + +// 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 that the middleware runs after the upstream call. +func (m *Middleware) Slot() middleware.Slot { return middleware.SlotOnResponse } + +// AcceptedContentTypes lists the response content types the middleware +// inspects. +func (m *Middleware) AcceptedContentTypes() []string { + return append([]string(nil), acceptedContentTypes...) +} + +// MetadataKeys returns the closed allowlist of keys this middleware +// may emit. +func (m *Middleware) MetadataKeys() []string { + return append([]string(nil), metadataKeys...) +} + +// MutationsSupported reports that this middleware never mutates the +// response. +func (m *Middleware) MutationsSupported() bool { return false } + +// Close releases any resources held by the middleware. The parser-set +// is stateless so this is a no-op. +func (m *Middleware) Close() error { return nil } + +// Invoke decodes the response body and emits token-usage and completion +// metadata. The decision is always DecisionAllow; parse errors degrade +// silently to omitted metadata rather than chain failures. +func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) { + out := &middleware.Output{Decision: middleware.DecisionAllow} + if in == nil { + return out, nil + } + + provider := lookupKV(in.Metadata, middleware.KeyLLMProvider) + if provider == "" { + return out, nil + } + + parser := m.parserByName(provider) + if parser == nil { + return out, nil + } + + // Upstreams compress the response when the client negotiated it + // (Claude Code sends Accept-Encoding: gzip). The transport leaves it + // compressed because the request carried an explicit Accept-Encoding, + // so the captured copy is gzip/deflate bytes — decompress it before + // parsing or token usage is silently lost. The forwarded client + // stream is untouched; this only affects our parse copy. + body := decodeResponseBody(in.RespBody, headerLookup(in.RespHeaders, "Content-Encoding")) + + contentType := headerLookup(in.RespHeaders, "Content-Type") + switch { + case isEventStream(contentType), isAWSEventStream(contentType): + out.Metadata = m.invokeStreaming(parser, body) + case isJSON(contentType): + out.Metadata = m.invokeBuffered(parser, in, contentType, body) + } + + return out, nil +} + +// invokeBuffered decodes a non-streaming JSON response body. Status +// codes >= 400 short-circuit because providers don't include usage on +// error responses. +func (m *Middleware) invokeBuffered(parser llm.Parser, in *middleware.Input, contentType string, body []byte) []middleware.KV { + if in.Status >= 400 { + return nil + } + + var md []middleware.KV + + usage, err := parser.ParseResponse(in.Status, contentType, body) + if err == nil { + md = appendUsage(md, usage) + } + + if completion := truncateCompletion(parser.ExtractCompletion(in.Status, contentType, body)); completion != "" && m.captureCompletion { + if m.redactPii { + completion = llm_guardrail.RedactPII(completion) + } + md = append(md, middleware.KV{Key: middleware.KeyLLMResponseCompletion, Value: completion}) + } + + return md +} + +// invokeStreaming walks the buffered SSE prefix and accumulates token +// deltas plus completion text. Truncated bodies are processed +// best-effort; partial usage is preferred over no metadata. +func (m *Middleware) invokeStreaming(parser llm.Parser, body []byte) []middleware.KV { + if len(body) == 0 { + return nil + } + + usage, completion := accumulateStream(parser.ProviderName(), body) + + var md []middleware.KV + if usage.InputTokens > 0 || usage.OutputTokens > 0 || usage.TotalTokens > 0 { + md = appendUsage(md, usage) + } + if c := truncateCompletion(completion); c != "" && m.captureCompletion { + if m.redactPii { + c = llm_guardrail.RedactPII(c) + } + md = append(md, middleware.KV{Key: middleware.KeyLLMResponseCompletion, Value: c}) + } + return md +} + +// parserByName returns the parser matching the provider label emitted +// by llm_request_parser, or nil when none claims it. +func (m *Middleware) parserByName(name string) llm.Parser { + for _, p := range m.parsers { + if p.ProviderName() == name { + return p + } + } + return nil +} + +// appendUsage emits the three baseline token-count metadata keys plus +// optional cached / cache-creation bucket counts when nonzero. Total +// is computed when the provider omitted one but reported per-direction +// counts; cache buckets are excluded from the legacy total because +// llm.input_tokens already absorbs the OpenAI cached subset and the +// sum-of-everything is a separate downstream concern. +func appendUsage(md []middleware.KV, usage llm.Usage) []middleware.KV { + total := usage.TotalTokens + if total == 0 && (usage.InputTokens > 0 || usage.OutputTokens > 0) { + total = usage.InputTokens + usage.OutputTokens + } + md = append(md, + middleware.KV{Key: middleware.KeyLLMInputTokens, Value: strconv.FormatInt(usage.InputTokens, 10)}, + middleware.KV{Key: middleware.KeyLLMOutputTokens, Value: strconv.FormatInt(usage.OutputTokens, 10)}, + middleware.KV{Key: middleware.KeyLLMTotalTokens, Value: strconv.FormatInt(total, 10)}, + ) + if usage.CachedInputTokens > 0 { + md = append(md, middleware.KV{ + Key: middleware.KeyLLMCachedInputTokens, + Value: strconv.FormatInt(usage.CachedInputTokens, 10), + }) + } + if usage.CacheCreationTokens > 0 { + md = append(md, middleware.KV{ + Key: middleware.KeyLLMCacheCreationTokens, + Value: strconv.FormatInt(usage.CacheCreationTokens, 10), + }) + } + return md +} + +// truncateCompletion clamps an extracted completion to maxCompletionBytes. +// The cut is rune-safe so we never split a multi-byte UTF-8 sequence. +func truncateCompletion(s string) string { + if len(s) <= maxCompletionBytes { + return s + } + cut := maxCompletionBytes + for cut > 0 && !utf8.RuneStart(s[cut]) { + cut-- + } + return s[:cut] +} + +func lookupKV(kvs []middleware.KV, key string) string { + for _, kv := range kvs { + if kv.Key == key { + return kv.Value + } + } + return "" +} + +func headerLookup(h []middleware.KV, name string) string { + lower := strings.ToLower(name) + for _, kv := range h { + if strings.ToLower(kv.Key) == lower { + return kv.Value + } + } + return "" +} + +func isEventStream(contentType string) bool { + return strings.Contains(strings.ToLower(contentType), "text/event-stream") +} + +// isAWSEventStream reports whether contentType is the AWS binary event-stream +// framing used by Bedrock's streaming endpoints. +func isAWSEventStream(contentType string) bool { + return strings.Contains(strings.ToLower(contentType), "application/vnd.amazon.eventstream") +} + +func isJSON(contentType string) bool { + lower := strings.ToLower(contentType) + return strings.Contains(lower, "application/json") || strings.Contains(lower, "+json") +} + +// decodeResponseBody returns body decompressed per its Content-Encoding, +// or the original bytes when the encoding is identity, unrecognised +// (e.g. br — no stdlib decoder), or the body isn't actually compressed. +// Decoding is best-effort: a truncated stream (capture hit the byte cap) +// yields the decompressed prefix rather than an error, which is enough to +// recover the leading message_start usage on Anthropic SSE. +func decodeResponseBody(body []byte, contentEncoding string) []byte { + enc := strings.ToLower(strings.TrimSpace(contentEncoding)) + // Content-Encoding may list multiple codings; the last applied is + // the outermost on the wire. + if idx := strings.LastIndex(enc, ","); idx >= 0 { + enc = strings.TrimSpace(enc[idx+1:]) + } + switch enc { + case "", "identity": + return body + case "gzip", "x-gzip": + zr, err := gzip.NewReader(bytes.NewReader(body)) + if err != nil { + return body + } + defer zr.Close() + if out := readCapped(zr); len(out) > 0 { + return out + } + return body + case "deflate": + // "deflate" on the wire is usually zlib-wrapped; fall back to raw + // flate when there's no zlib header. + if zr, err := zlib.NewReader(bytes.NewReader(body)); err == nil { + defer zr.Close() + if out := readCapped(zr); len(out) > 0 { + return out + } + return body + } + fr := flate.NewReader(bytes.NewReader(body)) + defer fr.Close() + if out := readCapped(fr); len(out) > 0 { + return out + } + return body + default: + return body + } +} + +// readCapped reads at most maxDecodedBytes from r, discarding any excess. +// Best-effort: a read error returns whatever was decoded so far, which is +// enough for the parser to recover leading usage events. +func readCapped(r io.Reader) []byte { + out, _ := io.ReadAll(io.LimitReader(r, maxDecodedBytes)) + return out +} diff --git a/proxy/internal/middleware/builtin/llm_response_parser/middleware_test.go b/proxy/internal/middleware/builtin/llm_response_parser/middleware_test.go new file mode 100644 index 000000000..084118802 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_response_parser/middleware_test.go @@ -0,0 +1,433 @@ +package llm_response_parser + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +func loadFixture(t *testing.T, name string) []byte { + t.Helper() + root, err := os.Getwd() + require.NoError(t, err, "must resolve cwd to locate fixture") + + dir := root + for i := 0; i < 8; i++ { + candidate := filepath.Join(dir, "proxy", "internal", "llm", "fixtures", name) + if data, err := os.ReadFile(candidate); err == nil { + return data + } + parent := filepath.Dir(dir) + if parent == dir { + break + } + dir = parent + } + t.Fatalf("fixture %q not found relative to %q", name, root) + return nil +} + +func metaValue(kvs []middleware.KV, key string) (string, bool) { + for _, kv := range kvs { + if kv.Key == key { + return kv.Value, true + } + } + return "", false +} + +func newTestMiddleware(t *testing.T) *Middleware { + t.Helper() + mw, err := Factory{}.New(nil) + require.NoError(t, err, "factory must accept empty config") + concrete, ok := mw.(*Middleware) + require.True(t, ok, "factory must return *Middleware") + return concrete +} + +func TestMiddleware_StaticSurface(t *testing.T) { + m := newTestMiddleware(t) + assert.Equal(t, ID, m.ID(), "ID must match registry constant") + assert.Equal(t, "1.0.0", m.Version(), "Version must be 1.0.0") + assert.Equal(t, middleware.SlotOnResponse, m.Slot(), "Slot must be SlotOnResponse") + assert.False(t, m.MutationsSupported(), "response parser does not mutate") + assert.ElementsMatch(t, + []string{"application/json", "text/event-stream"}, + m.AcceptedContentTypes(), + "AcceptedContentTypes must list JSON and SSE", + ) + assert.ElementsMatch(t, + []string{ + middleware.KeyLLMInputTokens, + middleware.KeyLLMOutputTokens, + middleware.KeyLLMTotalTokens, + middleware.KeyLLMCachedInputTokens, + middleware.KeyLLMCacheCreationTokens, + middleware.KeyLLMResponseCompletion, + }, + m.MetadataKeys(), + "MetadataKeys must be the documented response-side keys, including the optional cache buckets emitted only when nonzero", + ) + require.NoError(t, m.Close(), "Close must be a no-op") +} + +func TestFactory_AcceptsEmptyAndNullConfig(t *testing.T) { + for name, raw := range map[string][]byte{ + "nil": nil, + "empty": {}, + "null": []byte("null"), + "obj": []byte("{}"), + "ws": []byte(" "), + } { + t.Run(name, func(t *testing.T) { + mw, err := Factory{}.New(raw) + require.NoError(t, err, "factory must accept %s config", name) + require.NotNil(t, mw, "factory must return middleware for %s", name) + }) + } +} + +func TestFactory_RejectsMalformedJSON(t *testing.T) { + _, err := Factory{}.New([]byte("not-json")) + require.Error(t, err, "malformed config must surface a decode error") +} + +func TestInvoke_OpenAIBuffered(t *testing.T) { + m := newTestMiddleware(t) + body := loadFixture(t, "openai_chat_completion.json") + + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "application/json"}}, + RespBody: body, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o-mini"}, + }, + } + + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "Invoke must not error on a valid buffered response") + require.Equal(t, middleware.DecisionAllow, out.Decision, "decision must be Allow") + + in123, ok := metaValue(out.Metadata, middleware.KeyLLMInputTokens) + require.True(t, ok, "input tokens must be emitted") + assert.Equal(t, "123", in123, "input tokens must match fixture prompt_tokens") + + outTok, ok := metaValue(out.Metadata, middleware.KeyLLMOutputTokens) + require.True(t, ok, "output tokens must be emitted") + assert.Equal(t, "45", outTok, "output tokens must match fixture completion_tokens") + + totTok, ok := metaValue(out.Metadata, middleware.KeyLLMTotalTokens) + require.True(t, ok, "total tokens must be emitted") + assert.Equal(t, "168", totTok, "total tokens must match fixture") + + completion, ok := metaValue(out.Metadata, middleware.KeyLLMResponseCompletion) + require.True(t, ok, "completion must be emitted") + assert.Equal(t, "Hello, world!", completion, "completion text must match fixture") +} + +func TestInvoke_AnthropicBuffered(t *testing.T) { + m := newTestMiddleware(t) + body := loadFixture(t, "anthropic_messages.json") + + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "application/json"}}, + RespBody: body, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "anthropic"}, + {Key: middleware.KeyLLMModel, Value: "claude-sonnet-4-5"}, + }, + } + + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "Invoke must not error on a valid buffered response") + require.Equal(t, middleware.DecisionAllow, out.Decision, "decision must be Allow") + + in123, _ := metaValue(out.Metadata, middleware.KeyLLMInputTokens) + assert.Equal(t, "123", in123, "input tokens must match anthropic fixture") + + outTok, _ := metaValue(out.Metadata, middleware.KeyLLMOutputTokens) + assert.Equal(t, "45", outTok, "output tokens must match anthropic fixture") + + totTok, _ := metaValue(out.Metadata, middleware.KeyLLMTotalTokens) + assert.Equal(t, "168", totTok, "total tokens must be input+output for anthropic") + + completion, ok := metaValue(out.Metadata, middleware.KeyLLMResponseCompletion) + require.True(t, ok, "completion must be emitted for anthropic") + assert.Equal(t, "Hello, world!", completion, "completion text must match fixture") +} + +// TestInvoke_OpenAICachedTokensSurfaceOnMetadata covers the +// end-to-end path from the JSON usage block to the +// llm.cached_input_tokens metadata key the cost meter consumes. +// llm.cache_creation_tokens is NOT emitted for OpenAI because +// OpenAI has no cache_creation analogue. +func TestInvoke_OpenAICachedTokensSurfaceOnMetadata(t *testing.T) { + m := newTestMiddleware(t) + body := []byte(`{"usage":{"prompt_tokens":1024,"completion_tokens":200,"total_tokens":1224,"prompt_tokens_details":{"cached_tokens":768}}}`) + + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "application/json"}}, + RespBody: body, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + }, + } + + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err) + cached, ok := metaValue(out.Metadata, middleware.KeyLLMCachedInputTokens) + require.True(t, ok, "cached_input_tokens must land on the bag when the OpenAI response carries cached_tokens") + assert.Equal(t, "768", cached) + + _, hasCreation := metaValue(out.Metadata, middleware.KeyLLMCacheCreationTokens) + assert.False(t, hasCreation, "cache_creation_tokens must NOT be emitted for OpenAI — no analogue in the OpenAI shape") +} + +// TestInvoke_AnthropicCacheBucketsSurfaceOnMetadata covers the +// Anthropic shape: both cache_read and cache_creation values flow +// onto the metadata bag so the cost meter can apply per-bucket +// rates. +func TestInvoke_AnthropicCacheBucketsSurfaceOnMetadata(t *testing.T) { + m := newTestMiddleware(t) + body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768,"cache_creation_input_tokens":512}}`) + + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "application/json"}}, + RespBody: body, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "anthropic"}, + {Key: middleware.KeyLLMModel, Value: "claude-sonnet-4-5"}, + }, + } + + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err) + + cached, ok := metaValue(out.Metadata, middleware.KeyLLMCachedInputTokens) + require.True(t, ok, "cache_read_input_tokens lands under cached_input_tokens — same key carries OpenAI cached subset and Anthropic cache reads, meter switches formula on provider") + assert.Equal(t, "768", cached) + + creation, ok := metaValue(out.Metadata, middleware.KeyLLMCacheCreationTokens) + require.True(t, ok, "cache_creation_input_tokens lands under cache_creation_tokens for Anthropic") + assert.Equal(t, "512", creation) +} + +func TestInvoke_NoProviderMetadata_NoOp(t *testing.T) { + m := newTestMiddleware(t) + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "application/json"}}, + RespBody: loadFixture(t, "openai_chat_completion.json"), + } + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "missing provider metadata is not an error") + assert.Equal(t, middleware.DecisionAllow, out.Decision, "decision must be Allow") + assert.Empty(t, out.Metadata, "no metadata when provider context is missing") +} + +func TestInvoke_UnknownProvider_NoOp(t *testing.T) { + m := newTestMiddleware(t) + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "application/json"}}, + RespBody: loadFixture(t, "openai_chat_completion.json"), + Metadata: []middleware.KV{{Key: middleware.KeyLLMProvider, Value: "cohere"}}, + } + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "unknown provider must not surface an error") + assert.Empty(t, out.Metadata, "unknown providers emit no metadata") +} + +func TestInvoke_ErrorStatus_NoUsageEmitted(t *testing.T) { + m := newTestMiddleware(t) + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 500, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "application/json"}}, + RespBody: []byte(`{"error":{"message":"upstream blew up"}}`), + Metadata: []middleware.KV{{Key: middleware.KeyLLMProvider, Value: "openai"}}, + } + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "error responses must not surface as middleware error") + _, ok := metaValue(out.Metadata, middleware.KeyLLMInputTokens) + assert.False(t, ok, "no usage metadata on >=400 responses") +} + +func TestInvoke_NonInspectedContentType_NoOp(t *testing.T) { + m := newTestMiddleware(t) + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "text/plain"}}, + RespBody: []byte("not json"), + Metadata: []middleware.KV{{Key: middleware.KeyLLMProvider, Value: "openai"}}, + } + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "Invoke must tolerate non-inspected content types") + assert.Empty(t, out.Metadata, "no metadata for non-JSON, non-SSE bodies") +} + +func TestInvoke_NilInput(t *testing.T) { + m := newTestMiddleware(t) + out, err := m.Invoke(context.Background(), nil) + require.NoError(t, err, "nil input must not error") + require.Equal(t, middleware.DecisionAllow, out.Decision, "decision must be Allow even on nil input") + assert.Empty(t, out.Metadata, "no metadata for nil input") +} + +func TestInvoke_CompletionTruncatedAt3500Bytes(t *testing.T) { + m := newTestMiddleware(t) + long := strings.Repeat("x", 5000) + body := []byte(`{"id":"x","choices":[{"message":{"role":"assistant","content":"` + long + `"}}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`) + + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "application/json"}}, + RespBody: body, + Metadata: []middleware.KV{{Key: middleware.KeyLLMProvider, Value: "openai"}}, + } + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "long-completion body must parse cleanly") + + completion, ok := metaValue(out.Metadata, middleware.KeyLLMResponseCompletion) + require.True(t, ok, "completion must be emitted for long body") + assert.LessOrEqual(t, len(completion), maxCompletionBytes, "completion must be truncated to <=3500 bytes") + assert.Equal(t, maxCompletionBytes, len(completion), "completion must be truncated exactly at the cap when input is ASCII and longer") +} + +// TestInvoke_RedactPii_RedactsCompletionBeforeEmit covers the GC contract on +// the response leg: when the synthesiser sets redact_pii=true, the value +// emitted as llm.response_completion must already be redacted, so the +// access-log row never carries raw emails / SSNs / phones the model generated. +// Without this, the response side leaked dozens of raw PII tokens per request. +func TestInvoke_RedactPii_RedactsCompletionBeforeEmit(t *testing.T) { + mw, err := Factory{}.New([]byte(`{"redact_pii":true}`)) + require.NoError(t, err) + + piiCompletion := "Sample record: Alice Johnson, alice.johnson@example.com, SSN 123-45-6789, phone (202) 555-0147. Bob: 202/555/0108." + body := []byte(`{"id":"x","choices":[{"message":{"role":"assistant","content":"` + piiCompletion + `"}}],"usage":{"prompt_tokens":10,"completion_tokens":50,"total_tokens":60}}`) + + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "application/json"}}, + RespBody: body, + Metadata: []middleware.KV{{Key: middleware.KeyLLMProvider, Value: "openai"}}, + } + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + + completion, ok := metaValue(out.Metadata, middleware.KeyLLMResponseCompletion) + require.True(t, ok, "completion key must be emitted") + assert.Contains(t, completion, "[REDACTED:email]", "email must be redacted before emit") + assert.Contains(t, completion, "[REDACTED:ssn]", "ssn must be redacted before emit") + assert.Contains(t, completion, "[REDACTED:phone]", "phone must be redacted before emit") + assert.NotContains(t, completion, "alice.johnson@example.com", "raw email must not survive") + assert.NotContains(t, completion, "123-45-6789", "raw SSN must not survive") + assert.NotContains(t, completion, "(202) 555-0147", "parens-phone must not survive") + assert.NotContains(t, completion, "202/555/0108", "slash-phone must not survive") +} + +// TestInvoke_CaptureCompletionOff_DoesNotEmitCompletion mirrors the request +// parser test: when capture_completion=false (operator has enable_prompt_ +// collection off), llm.response_completion MUST NOT appear in the access log. +// The token / cost / usage facts the response parser also emits stay so +// operators still get billing data on log-only mode. +func TestInvoke_CaptureCompletionOff_DoesNotEmitCompletion(t *testing.T) { + mw, err := Factory{}.New([]byte(`{"capture_completion":false}`)) + require.NoError(t, err) + body := []byte(`{"id":"x","choices":[{"message":{"role":"assistant","content":"alice@example.com 123-45-6789"}}],"usage":{"prompt_tokens":10,"completion_tokens":20,"total_tokens":30}}`) + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "application/json"}}, + RespBody: body, + Metadata: []middleware.KV{{Key: middleware.KeyLLMProvider, Value: "openai"}}, + } + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + + _, ok := metaValue(out.Metadata, middleware.KeyLLMResponseCompletion) + assert.False(t, ok, "llm.response_completion must NOT be emitted when capture_completion is false") + + // Token facts must still flow. + _, ok = metaValue(out.Metadata, middleware.KeyLLMInputTokens) + assert.True(t, ok, "input tokens fact must still be emitted") + _, ok = metaValue(out.Metadata, middleware.KeyLLMOutputTokens) + assert.True(t, ok, "output tokens fact must still be emitted") +} + +// TestInvoke_CaptureCompletionUnset_PreservesLegacyEmission documents the +// default behavior: empty config keeps emitting completion, so callers +// without the toggle aren't broken. +func TestInvoke_CaptureCompletionUnset_PreservesLegacyEmission(t *testing.T) { + mw, err := Factory{}.New([]byte(`{}`)) + require.NoError(t, err) + body := []byte(`{"id":"x","choices":[{"message":{"role":"assistant","content":"hello"}}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`) + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "application/json"}}, + RespBody: body, + Metadata: []middleware.KV{{Key: middleware.KeyLLMProvider, Value: "openai"}}, + } + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + _, ok := metaValue(out.Metadata, middleware.KeyLLMResponseCompletion) + assert.True(t, ok, "absent capture_completion must preserve emission (backwards-compatible default)") +} + +// TestInvoke_RedactPii_OffShipsRawCompletion covers the inverse: with +// redact_pii=false (default) the model output is shipped verbatim. +func TestInvoke_RedactPii_OffShipsRawCompletion(t *testing.T) { + mw, err := Factory{}.New(nil) + require.NoError(t, err) + + body := []byte(`{"id":"x","choices":[{"message":{"role":"assistant","content":"alice@example.com 123-45-6789"}}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`) + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "application/json"}}, + RespBody: body, + Metadata: []middleware.KV{{Key: middleware.KeyLLMProvider, Value: "openai"}}, + } + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + + completion, ok := metaValue(out.Metadata, middleware.KeyLLMResponseCompletion) + require.True(t, ok) + assert.Contains(t, completion, "alice@example.com", "redact off → raw email passes through") + assert.Contains(t, completion, "123-45-6789", "redact off → raw SSN passes through") + assert.NotContains(t, completion, "[REDACTED:", "redact off → no markers") +} + +func TestInvoke_CompletionTruncationRuneSafe(t *testing.T) { + rune4 := "\xf0\x9f\x98\x80" // 4-byte emoji + body := strings.Repeat("a", maxCompletionBytes-1) + rune4 + require.Greater(t, len(body), maxCompletionBytes, "test setup must exceed the cap") + + got := truncateCompletion(body) + assert.True(t, len(got) < maxCompletionBytes, "truncated bytes must drop the partial rune entirely") + assert.NotContains(t, got, "\x80", "truncated text must not end on a continuation byte") +} diff --git a/proxy/internal/middleware/builtin/llm_response_parser/responses_stream_test.go b/proxy/internal/middleware/builtin/llm_response_parser/responses_stream_test.go new file mode 100644 index 000000000..0475b8cca --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_response_parser/responses_stream_test.go @@ -0,0 +1,69 @@ +package llm_response_parser + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +// TestInvoke_OpenAIResponsesStreaming is the regression guard for the live +// bug where Codex hits /v1/responses (the OpenAI Responses API), whose SSE +// shape differs from chat.completions: completion text rides +// response.output_text.delta and usage rides response.completed under +// response.usage. The old parser only knew the chat.completions shape, so +// resp_meta came back empty (no tokens, no cost). +func TestInvoke_OpenAIResponsesStreaming(t *testing.T) { + m := newTestMiddleware(t) + body := loadFixture(t, "openai_responses_stream.txt") + + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "text/event-stream; charset=utf-8"}}, + RespBody: body, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-5.5"}, + }, + } + + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "Invoke must not error on a Responses-API streaming body") + + inTok, ok := metaValue(out.Metadata, middleware.KeyLLMInputTokens) + require.True(t, ok, "input tokens must be emitted from a Responses-API stream") + assert.Equal(t, "123", inTok, "input_tokens must come from response.completed usage") + + outTok, _ := metaValue(out.Metadata, middleware.KeyLLMOutputTokens) + assert.Equal(t, "45", outTok, "output_tokens must come from response.completed usage") + + totTok, _ := metaValue(out.Metadata, middleware.KeyLLMTotalTokens) + assert.Equal(t, "168", totTok, "total_tokens must come from response.completed usage") + + cached, ok := metaValue(out.Metadata, middleware.KeyLLMCachedInputTokens) + require.True(t, ok, "cached input tokens must surface from input_tokens_details") + assert.Equal(t, "40", cached, "cached_tokens subset must surface for cost discounting") + + completion, ok := metaValue(out.Metadata, middleware.KeyLLMResponseCompletion) + require.True(t, ok, "completion must be emitted for Responses-API streams") + assert.Equal(t, "Hello, world!", completion, "output_text.delta events must concatenate") +} + +// TestAccumulateOpenAIStream_ResponsesNoUsage confirms that a Responses-API +// stream with text but no terminal usage frame still yields the completion +// and leaves tokens at zero rather than erroring. +func TestAccumulateOpenAIStream_ResponsesNoUsage(t *testing.T) { + body := []byte(`event: response.output_text.delta +data: {"type":"response.output_text.delta","delta":"partial"} + +`) + + usage, completion := accumulateOpenAIStream(body) + assert.Equal(t, int64(0), usage.InputTokens, "no usage frame leaves input tokens at zero") + assert.Equal(t, int64(0), usage.OutputTokens, "no usage frame leaves output tokens at zero") + assert.Equal(t, "partial", completion, "output_text deltas accumulate even without a usage frame") +} diff --git a/proxy/internal/middleware/builtin/llm_response_parser/streaming.go b/proxy/internal/middleware/builtin/llm_response_parser/streaming.go new file mode 100644 index 000000000..9d3e8a204 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_response_parser/streaming.go @@ -0,0 +1,269 @@ +package llm_response_parser + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "strings" + + "github.com/netbirdio/netbird/proxy/internal/llm" +) + +// openAIDoneSentinel is the OpenAI end-of-stream marker. The scanner +// stops once this data frame is observed. +const openAIDoneSentinel = "[DONE]" + +// accumulateStream walks the SSE byte slice, dispatches per provider, +// and returns the running token-usage and concatenated completion text. +// Errors from the scanner short-circuit accumulation but never panic +// — partial results are returned for truncated bodies. +func accumulateStream(provider string, body []byte) (llm.Usage, string) { + switch provider { + case "openai": + return accumulateOpenAIStream(body) + case "anthropic": + return accumulateAnthropicStream(body) + case llm.ProviderNameBedrock: + return accumulateBedrockStream(body) + default: + return llm.Usage{}, "" + } +} + +// openAIStreamUsage is the usage block shared by both OpenAI streaming +// envelopes. Pointer fields tell "absent" from zero; the chat.completions +// (prompt_/completion_) and Responses-API (input_/output_) names are both +// accepted so a single decode covers either endpoint. +type openAIStreamUsage struct { + PromptTokens *int64 `json:"prompt_tokens"` + CompletionTokens *int64 `json:"completion_tokens"` + InputTokens *int64 `json:"input_tokens"` + OutputTokens *int64 `json:"output_tokens"` + TotalTokens *int64 `json:"total_tokens"` + PromptTokensDetails *struct { + CachedTokens *int64 `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + InputTokensDetails *struct { + CachedTokens *int64 `json:"cached_tokens"` + } `json:"input_tokens_details"` +} + +// openAIStreamChunk matches both OpenAI streaming envelopes. The +// chat.completions chunk carries text in choices[].delta.content and a +// trailing top-level usage block. The Responses API (/v1/responses) emits +// typed events instead: completion text rides response.output_text.delta +// (top-level "delta" string) and the final usage rides response.completed +// under response.usage. Only fields used for accumulation are declared. +type openAIStreamChunk struct { + Choices []struct { + Delta struct { + Content string `json:"content"` + } `json:"delta"` + } `json:"choices"` + Usage *openAIStreamUsage `json:"usage"` + + Type string `json:"type"` + Delta json.RawMessage `json:"delta"` + Response *struct { + Usage *openAIStreamUsage `json:"usage"` + } `json:"response"` +} + +// accumulateOpenAIStream sums per-chunk content deltas and lifts the usage +// block off the final frame, handling both the chat.completions and the +// Responses-API event shapes. Clients without stream_options.include_usage +// (chat.completions) and any provider that omits the final usage simply +// leave tokens at zero; the caller chooses what to emit. +func accumulateOpenAIStream(body []byte) (llm.Usage, string) { + var ( + usage llm.Usage + completion strings.Builder + ) + scanner := llm.NewScanner(bytes.NewReader(body)) + for { + ev, err := scanner.Next() + if err != nil { + if errors.Is(err, io.EOF) { + break + } + break + } + if ev.Data == "" || ev.Data == openAIDoneSentinel { + if ev.Data == openAIDoneSentinel { + break + } + continue + } + + var chunk openAIStreamChunk + if err := json.Unmarshal([]byte(ev.Data), &chunk); err != nil { + continue + } + for _, c := range chunk.Choices { + completion.WriteString(c.Delta.Content) + } + if chunk.Type == "response.output_text.delta" { + if s, ok := decodeJSONString(chunk.Delta); ok { + completion.WriteString(s) + } + } + + u := chunk.Usage + if u == nil && chunk.Response != nil { + u = chunk.Response.Usage + } + if u != nil { + usage.InputTokens = pickInt64(u.InputTokens, u.PromptTokens) + usage.OutputTokens = pickInt64(u.OutputTokens, u.CompletionTokens) + usage.TotalTokens = derefInt64(u.TotalTokens) + if u.InputTokensDetails != nil { + if v := derefInt64(u.InputTokensDetails.CachedTokens); v > 0 { + usage.CachedInputTokens = v + } + } + if usage.CachedInputTokens == 0 && u.PromptTokensDetails != nil { + usage.CachedInputTokens = derefInt64(u.PromptTokensDetails.CachedTokens) + } + if usage.TotalTokens == 0 && (usage.InputTokens > 0 || usage.OutputTokens > 0) { + usage.TotalTokens = usage.InputTokens + usage.OutputTokens + } + } + } + return usage, completion.String() +} + +// decodeJSONString unmarshals a JSON-encoded string value, returning +// ok=false when the raw message is empty or not a string. +func decodeJSONString(raw json.RawMessage) (string, bool) { + if len(raw) == 0 { + return "", false + } + var s string + if err := json.Unmarshal(raw, &s); err != nil { + return "", false + } + return s, true +} + +// anthropicStreamEvent captures the union of Messages-API stream event +// payloads we care about. Each named event on the wire fills only its +// shape's fields; unknown keys are ignored. +type anthropicStreamEvent struct { + Type string `json:"type"` + Message *struct { + Usage *struct { + InputTokens *int64 `json:"input_tokens"` + OutputTokens *int64 `json:"output_tokens"` + CacheReadInputTokens *int64 `json:"cache_read_input_tokens"` + CacheCreationInputTokens *int64 `json:"cache_creation_input_tokens"` + } `json:"usage"` + } `json:"message"` + Delta *struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"delta"` + Usage *struct { + InputTokens *int64 `json:"input_tokens"` + OutputTokens *int64 `json:"output_tokens"` + CacheReadInputTokens *int64 `json:"cache_read_input_tokens"` + CacheCreationInputTokens *int64 `json:"cache_creation_input_tokens"` + } `json:"usage"` +} + +// accumulateAnthropicStream tracks input_tokens from message_start, +// output_tokens from message_delta, and concatenates text_delta payloads +// from content_block_delta events. Final usage prefers message_delta +// values which carry the post-completion totals. +func accumulateAnthropicStream(body []byte) (llm.Usage, string) { + var ( + usage llm.Usage + completion strings.Builder + ) + scanner := llm.NewScanner(bytes.NewReader(body)) + for { + ev, err := scanner.Next() + if err != nil { + if errors.Is(err, io.EOF) { + break + } + break + } + if ev.Data == "" { + continue + } + + var payload anthropicStreamEvent + if err := json.Unmarshal([]byte(ev.Data), &payload); err != nil { + continue + } + + eventType := ev.Type + if eventType == "" { + eventType = payload.Type + } + applyAnthropicStreamEvent(eventType, payload, &usage, &completion) + } + if usage.InputTokens > 0 || usage.OutputTokens > 0 { + usage.TotalTokens = usage.InputTokens + usage.OutputTokens + usage.CachedInputTokens + usage.CacheCreationTokens + } + return usage, completion.String() +} + +// applyAnthropicStreamEvent folds one parsed Anthropic Messages stream event +// into the running usage/completion. Shared by the SSE accumulator and the +// Bedrock InvokeModel event-stream, whose chunks wrap the same event JSON. +func applyAnthropicStreamEvent(eventType string, payload anthropicStreamEvent, usage *llm.Usage, completion *strings.Builder) { + switch eventType { + case "message_start": + if payload.Message != nil && payload.Message.Usage != nil { + if v := derefInt64(payload.Message.Usage.InputTokens); v > 0 { + usage.InputTokens = v + } + if v := derefInt64(payload.Message.Usage.OutputTokens); v > 0 { + usage.OutputTokens = v + } + if v := derefInt64(payload.Message.Usage.CacheReadInputTokens); v > 0 { + usage.CachedInputTokens = v + } + if v := derefInt64(payload.Message.Usage.CacheCreationInputTokens); v > 0 { + usage.CacheCreationTokens = v + } + } + case "content_block_delta": + if payload.Delta != nil && payload.Delta.Type == "text_delta" { + completion.WriteString(payload.Delta.Text) + } + case "message_delta": + if payload.Usage != nil { + if v := derefInt64(payload.Usage.InputTokens); v > 0 { + usage.InputTokens = v + } + if v := derefInt64(payload.Usage.OutputTokens); v > 0 { + usage.OutputTokens = v + } + if v := derefInt64(payload.Usage.CacheReadInputTokens); v > 0 { + usage.CachedInputTokens = v + } + if v := derefInt64(payload.Usage.CacheCreationInputTokens); v > 0 { + usage.CacheCreationTokens = v + } + } + case "message_stop": + // No-op; Anthropic does not emit usage here. + } +} + +func pickInt64(preferred, fallback *int64) int64 { + if preferred != nil { + return *preferred + } + return derefInt64(fallback) +} + +func derefInt64(v *int64) int64 { + if v == nil { + return 0 + } + return *v +} diff --git a/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock.go b/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock.go new file mode 100644 index 000000000..a82a9cdbc --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock.go @@ -0,0 +1,110 @@ +package llm_response_parser + +import ( + "bytes" + "encoding/json" + "strings" + + "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream" + + "github.com/netbirdio/netbird/proxy/internal/llm" +) + +// bedrockEventTypeHeader names each AWS event-stream frame's event type. +const bedrockEventTypeHeader = ":event-type" + +// accumulateBedrockStream decodes the AWS binary event-stream returned by +// Bedrock's streaming endpoints and folds it into running usage/completion. +// Two framings are handled: +// - InvokeModel (invoke-with-response-stream): each "chunk" frame's payload is +// {"bytes":""} wrapping a vendor-native (Anthropic) stream event. +// - Converse (converse-stream): native frames (contentBlockDelta, metadata, …) +// whose payload JSON carries text deltas and a final usage block. +// +// A truncated stream (cut at the capture cap) decodes best-effort: frames up to +// the cut are applied and the partial usage is returned. +func accumulateBedrockStream(body []byte) (llm.Usage, string) { + var ( + usage llm.Usage + completion strings.Builder + ) + dec := eventstream.NewDecoder() + r := bytes.NewReader(body) + for { + msg, err := dec.Decode(r, nil) + if err != nil { + break // EOF or a partial trailing frame — return what we have. + } + eventType := "" + if v := msg.Headers.Get(bedrockEventTypeHeader); v != nil { + eventType = v.String() + } + if eventType == "chunk" { + applyBedrockInvokeChunk(msg.Payload, &usage, &completion) + continue + } + applyConverseStreamEvent(eventType, msg.Payload, &usage, &completion) + } + if usage.TotalTokens == 0 && (usage.InputTokens > 0 || usage.OutputTokens > 0) { + usage.TotalTokens = usage.InputTokens + usage.OutputTokens + usage.CachedInputTokens + usage.CacheCreationTokens + } + return usage, completion.String() +} + +// applyBedrockInvokeChunk decodes an InvokeModel stream "chunk" frame +// ({"bytes":""}) and folds the wrapped Anthropic event +// into usage/completion via the shared accumulator. +func applyBedrockInvokeChunk(payload []byte, usage *llm.Usage, completion *strings.Builder) { + var wrap struct { + Bytes []byte `json:"bytes"` // base64 string — encoding/json decodes it + } + if err := json.Unmarshal(payload, &wrap); err != nil || len(wrap.Bytes) == 0 { + return + } + var ev anthropicStreamEvent + if err := json.Unmarshal(wrap.Bytes, &ev); err != nil { + return + } + applyAnthropicStreamEvent(ev.Type, ev, usage, completion) +} + +// converseStreamEvent captures the Converse stream frames carrying completion +// text (contentBlockDelta) and the final token usage (metadata). +type converseStreamEvent struct { + Delta *struct { + Text string `json:"text"` + } `json:"delta"` + Usage *struct { + InputTokens int64 `json:"inputTokens"` + OutputTokens int64 `json:"outputTokens"` + TotalTokens int64 `json:"totalTokens"` + } `json:"usage"` +} + +// applyConverseStreamEvent folds one native Converse stream frame into the +// running usage/completion: contentBlockDelta carries assistant text, and the +// trailing metadata frame carries the final usage block. +func applyConverseStreamEvent(eventType string, payload []byte, usage *llm.Usage, completion *strings.Builder) { + var ev converseStreamEvent + if err := json.Unmarshal(payload, &ev); err != nil { + return + } + switch eventType { + case "contentBlockDelta": + if ev.Delta != nil { + completion.WriteString(ev.Delta.Text) + } + case "metadata": + if ev.Usage != nil { + if ev.Usage.InputTokens > 0 { + usage.InputTokens = ev.Usage.InputTokens + } + if ev.Usage.OutputTokens > 0 { + usage.OutputTokens = ev.Usage.OutputTokens + } + if ev.Usage.TotalTokens > 0 { + usage.TotalTokens = ev.Usage.TotalTokens + } + } + } +} diff --git a/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock_test.go b/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock_test.go new file mode 100644 index 000000000..f93505882 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock_test.go @@ -0,0 +1,74 @@ +package llm_response_parser + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream" + "github.com/stretchr/testify/require" +) + +// bedrockFrame encodes a single AWS event-stream frame with the given +// :event-type header and JSON payload, mirroring what Bedrock sends. +func bedrockFrame(t *testing.T, eventType string, payload []byte) []byte { + t.Helper() + var buf bytes.Buffer + enc := eventstream.NewEncoder() + err := enc.Encode(&buf, eventstream.Message{ + Headers: eventstream.Headers{{Name: ":event-type", Value: eventstream.StringValue(eventType)}}, + Payload: payload, + }) + require.NoError(t, err, "encode event-stream frame") + return buf.Bytes() +} + +func mustJSON(t *testing.T, v any) []byte { + t.Helper() + b, err := json.Marshal(v) + require.NoError(t, err) + return b +} + +func TestAccumulateBedrockStream_Invoke(t *testing.T) { + // invoke-with-response-stream: each "chunk" frame wraps a base64-encoded + // Anthropic stream event under {"bytes": ...}. + events := [][]byte{ + mustJSON(t, map[string]any{"type": "message_start", "message": map[string]any{"usage": map[string]any{"input_tokens": 13}}}), + mustJSON(t, map[string]any{"type": "content_block_delta", "delta": map[string]any{"type": "text_delta", "text": "po"}}), + mustJSON(t, map[string]any{"type": "content_block_delta", "delta": map[string]any{"type": "text_delta", "text": "ng"}}), + mustJSON(t, map[string]any{"type": "message_delta", "usage": map[string]any{"output_tokens": 5}}), + } + var body bytes.Buffer + for _, ev := range events { + wrap := mustJSON(t, map[string]any{"bytes": base64.StdEncoding.EncodeToString(ev)}) + body.Write(bedrockFrame(t, "chunk", wrap)) + } + + usage, completion := accumulateBedrockStream(body.Bytes()) + require.Equal(t, int64(13), usage.InputTokens, "input tokens from message_start") + require.Equal(t, int64(5), usage.OutputTokens, "output tokens from message_delta") + require.Equal(t, int64(18), usage.TotalTokens, "total is additive") + require.Equal(t, "pong", completion, "text deltas concatenated") +} + +func TestAccumulateBedrockStream_Converse(t *testing.T) { + var body bytes.Buffer + body.Write(bedrockFrame(t, "contentBlockDelta", mustJSON(t, map[string]any{"delta": map[string]any{"text": "po"}}))) + body.Write(bedrockFrame(t, "contentBlockDelta", mustJSON(t, map[string]any{"delta": map[string]any{"text": "ng"}}))) + body.Write(bedrockFrame(t, "metadata", mustJSON(t, map[string]any{"usage": map[string]any{"inputTokens": 11, "outputTokens": 3, "totalTokens": 14}}))) + + usage, completion := accumulateBedrockStream(body.Bytes()) + require.Equal(t, int64(11), usage.InputTokens, "input tokens from metadata frame") + require.Equal(t, int64(3), usage.OutputTokens, "output tokens from metadata frame") + require.Equal(t, int64(14), usage.TotalTokens, "total from metadata frame") + require.Equal(t, "pong", completion, "converse text deltas concatenated") +} + +func TestAccumulateBedrockStream_Truncated(t *testing.T) { + // A body cut mid-frame must not panic; partial usage is returned. + full := bedrockFrame(t, "metadata", mustJSON(t, map[string]any{"usage": map[string]any{"inputTokens": 11, "outputTokens": 3}})) + usage, _ := accumulateBedrockStream(full[:len(full)-4]) + require.Zero(t, usage.OutputTokens, "truncated trailing frame is dropped, not panicked on") +} diff --git a/proxy/internal/middleware/builtin/llm_response_parser/streaming_test.go b/proxy/internal/middleware/builtin/llm_response_parser/streaming_test.go new file mode 100644 index 000000000..400aac0bd --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_response_parser/streaming_test.go @@ -0,0 +1,169 @@ +package llm_response_parser + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +func TestInvoke_OpenAIStreamingWithUsage(t *testing.T) { + m := newTestMiddleware(t) + body := loadFixture(t, "openai_stream.txt") + + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "text/event-stream"}}, + RespBody: body, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "openai"}, + {Key: middleware.KeyLLMModel, Value: "gpt-4o-mini"}, + }, + } + + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "Invoke must not error on streaming OpenAI body") + + in123, _ := metaValue(out.Metadata, middleware.KeyLLMInputTokens) + assert.Equal(t, "123", in123, "input tokens must come from final-chunk usage block") + + outTok, _ := metaValue(out.Metadata, middleware.KeyLLMOutputTokens) + assert.Equal(t, "45", outTok, "output tokens must come from final-chunk usage block") + + totTok, _ := metaValue(out.Metadata, middleware.KeyLLMTotalTokens) + assert.Equal(t, "168", totTok, "total tokens must come from final-chunk usage block") + + completion, ok := metaValue(out.Metadata, middleware.KeyLLMResponseCompletion) + require.True(t, ok, "completion must be emitted for streaming responses") + assert.Equal(t, "Hello, world!", completion, "deltas must concatenate into the buffered fixture's text") +} + +func TestInvoke_OpenAIStreamingWithoutUsage(t *testing.T) { + body := []byte(`data: {"choices":[{"delta":{"content":"Hi"}}]} + +data: {"choices":[{"delta":{"content":" there"}}]} + +data: [DONE] + +`) + + usage, completion := accumulateOpenAIStream(body) + assert.Equal(t, int64(0), usage.InputTokens, "input tokens must stay zero without a usage frame") + assert.Equal(t, int64(0), usage.OutputTokens, "output tokens must stay zero without a usage frame") + assert.Equal(t, int64(0), usage.TotalTokens, "total tokens must stay zero without a usage frame") + assert.Equal(t, "Hi there", completion, "deltas must still accumulate when usage is absent") +} + +func TestInvoke_OpenAIStreamingNoUsage_OmitsUsageMetadata(t *testing.T) { + m := newTestMiddleware(t) + body := []byte(`data: {"choices":[{"delta":{"content":"Hello"}}]} + +data: [DONE] + +`) + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "text/event-stream"}}, + RespBody: body, + Metadata: []middleware.KV{{Key: middleware.KeyLLMProvider, Value: "openai"}}, + } + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "Invoke must not error on usage-less streams") + + _, hasIn := metaValue(out.Metadata, middleware.KeyLLMInputTokens) + _, hasOut := metaValue(out.Metadata, middleware.KeyLLMOutputTokens) + _, hasTot := metaValue(out.Metadata, middleware.KeyLLMTotalTokens) + assert.False(t, hasIn, "input tokens omitted when no usage frame") + assert.False(t, hasOut, "output tokens omitted when no usage frame") + assert.False(t, hasTot, "total tokens omitted when no usage frame") + + completion, ok := metaValue(out.Metadata, middleware.KeyLLMResponseCompletion) + require.True(t, ok, "completion must still be emitted from deltas") + assert.Equal(t, "Hello", completion, "completion must come from delta accumulation") +} + +func TestInvoke_AnthropicStreaming(t *testing.T) { + m := newTestMiddleware(t) + body := loadFixture(t, "anthropic_stream.txt") + + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "text/event-stream"}}, + RespBody: body, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: "anthropic"}, + {Key: middleware.KeyLLMModel, Value: "claude-sonnet-4-5"}, + }, + } + + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "Invoke must not error on streaming Anthropic body") + + in123, _ := metaValue(out.Metadata, middleware.KeyLLMInputTokens) + assert.Equal(t, "123", in123, "input tokens must come from message_start usage") + + outTok, _ := metaValue(out.Metadata, middleware.KeyLLMOutputTokens) + assert.Equal(t, "45", outTok, "output tokens must come from message_delta usage") + + totTok, _ := metaValue(out.Metadata, middleware.KeyLLMTotalTokens) + assert.Equal(t, "168", totTok, "total tokens must be input+output for anthropic streaming") + + completion, ok := metaValue(out.Metadata, middleware.KeyLLMResponseCompletion) + require.True(t, ok, "completion must be emitted from text_delta accumulation") + assert.Equal(t, "Hello, world!", completion, "anthropic streaming text must accumulate across content_block_delta events") +} + +func TestInvoke_StreamingTruncatedBody_BestEffort(t *testing.T) { + m := newTestMiddleware(t) + full := loadFixture(t, "anthropic_stream.txt") + cut := len(full) / 2 + truncated := full[:cut] + + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "text/event-stream"}}, + RespBody: truncated, + RespBodyTruncated: true, + Metadata: []middleware.KV{{Key: middleware.KeyLLMProvider, Value: "anthropic"}}, + } + + require.NotPanics(t, func() { + _, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "truncated streaming body must not surface as error") + }, "Invoke must never panic on a truncated SSE body") +} + +func TestInvoke_StreamingEmptyBody(t *testing.T) { + m := newTestMiddleware(t) + in := &middleware.Input{ + Slot: middleware.SlotOnResponse, + Status: 200, + RespHeaders: []middleware.KV{{Key: "Content-Type", Value: "text/event-stream"}}, + RespBody: nil, + Metadata: []middleware.KV{{Key: middleware.KeyLLMProvider, Value: "openai"}}, + } + out, err := m.Invoke(context.Background(), in) + require.NoError(t, err, "empty SSE body must not surface as error") + assert.Empty(t, out.Metadata, "no metadata for empty SSE body") +} + +func TestAccumulateAnthropicStream_PartialUsage(t *testing.T) { + body := []byte(`event: message_start +data: {"type":"message_start","message":{"usage":{"input_tokens":10}}} + +event: content_block_delta +data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"hi"}} + +`) + usage, completion := accumulateAnthropicStream(body) + assert.Equal(t, int64(10), usage.InputTokens, "partial input_tokens must survive truncated stream") + assert.Equal(t, int64(0), usage.OutputTokens, "output_tokens stays zero without message_delta") + assert.Equal(t, "hi", completion, "completion must come from observed text_delta events") +} diff --git a/proxy/internal/middleware/builtin/llm_router/factory.go b/proxy/internal/middleware/builtin/llm_router/factory.go new file mode 100644 index 000000000..3c3b607ac --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_router/factory.go @@ -0,0 +1,106 @@ +package llm_router + +import ( + "bytes" + "encoding/json" + "fmt" + "strings" + + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/proxy/internal/middleware/builtin" +) + +// ProviderRoute describes one upstream LLM provider the router can +// hand a request to. Models lists the model identifiers the provider +// claims; UpstreamScheme + UpstreamHost replace the synth target's +// placeholder URL on a match. UpstreamPath is the path component of +// the configured upstream URL — the router uses it to disambiguate +// providers that claim the same model: when more than one provider +// matches the model, the route whose UpstreamPath is a prefix of the +// incoming request path is preferred (longest match wins, empty path +// is the catchall). AuthHeaderName + AuthHeaderValue are the +// per-provider credential the router injects after stripping the +// vendor auth headers from the inbound request. +// +// AllowedGroupIDs is the union of source-group IDs across every +// enabled policy that authorises this provider. The router treats it +// as a hard filter: a route whose AllowedGroupIDs has no intersection +// with the caller's UserGroups is removed from the candidate list +// before the path-prefix tiebreak. A route with empty AllowedGroupIDs +// is unreachable; the synthesiser only emits policy-bound routes. +type ProviderRoute struct { + ID string `json:"id"` + // Vendor is the parser surface this provider speaks ("openai", + // "anthropic", …), matching the llm.provider value llm_request_parser + // emits from the request. When set, the router keeps a vendor-tagged + // request on a same-vendor route so catch-all gateways of a different + // vendor can't swallow it. Empty disables vendor filtering for this + // route. + Vendor string `json:"vendor,omitempty"` + Models []string `json:"models"` + UpstreamScheme string `json:"upstream_scheme"` + UpstreamHost string `json:"upstream_host"` + UpstreamPath string `json:"upstream_path,omitempty"` + AuthHeaderName string `json:"auth_header_name"` + AuthHeaderValue string `json:"auth_header_value"` + AllowedGroupIDs []string `json:"allowed_group_ids"` + // Vertex marks a Google Vertex AI provider. Vertex requests carry the + // model in the URL path, so the router selects this route by path + // (isVertexPath) and bypasses the model/vendor table entirely. + Vertex bool `json:"vertex,omitempty"` + // Bedrock marks an AWS Bedrock provider. Bedrock requests carry the model + // in the URL path (/model/{id}/{action}), so the router selects this route + // by path (isBedrockPath) and bypasses the model/vendor table; auth is the + // static AuthHeaderValue bearer token (no token minting). + Bedrock bool `json:"bedrock,omitempty"` + // GCPServiceAccountKeyB64 is a base64-encoded GCP service-account JSON + // key. When set, the router mints + refreshes a short-lived OAuth2 access + // token from it at request time and injects it as the auth header value + // (instead of the static AuthHeaderValue) — so the gateway holds a durable + // Vertex credential rather than a 1-hour token. + GCPServiceAccountKeyB64 string `json:"gcp_sa_key_b64,omitempty"` +} + +// Config is the on-wire configuration accepted by the factory. An +// empty Providers slice yields a router that denies every request as +// not-routable; the synthesiser is responsible for stamping the +// account's enabled providers into this slice. +type Config struct { + Providers []ProviderRoute `json:"providers"` +} + +// Factory builds llm_router 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 +// yield a router with an empty Providers slice — every request denies +// with model_not_routable. Non-empty payloads must parse cleanly so +// misconfigurations surface at chain build time. +func (Factory) New(rawConfig []byte) (middleware.Middleware, error) { + cfg := Config{} + if !isEmptyJSON(rawConfig) { + if err := json.Unmarshal(rawConfig, &cfg); err != nil { + return nil, fmt.Errorf("decode config: %w", err) + } + } + return New(cfg), nil +} + +// isEmptyJSON reports whether the payload is whitespace, null, or an +// empty object/array. The caller skips Unmarshal in that case so the +// zero-value Config flows through unchanged. +func isEmptyJSON(raw []byte) bool { + trimmed := strings.TrimSpace(string(bytes.TrimSpace(raw))) + switch trimmed { + case "", "null", "{}", "[]": + return true + } + return false +} + +func init() { + builtin.Register(Factory{}) +} diff --git a/proxy/internal/middleware/builtin/llm_router/middleware.go b/proxy/internal/middleware/builtin/llm_router/middleware.go new file mode 100644 index 000000000..73cc59c95 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_router/middleware.go @@ -0,0 +1,793 @@ +// Package llm_router implements the SlotOnRequest middleware that +// routes a request to an upstream LLM provider based on the model name +// emitted upstream by llm_request_parser. The router rewrites the +// request's outbound target (scheme + host), strips known LLM-vendor +// auth headers, and injects the per-provider auth header from the +// matched route. Unknown or unconfigured models deny with a 403 and +// the canonical llm_policy.model_not_routable code. +package llm_router + +import ( + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "fmt" + "net/http" + "net/url" + "sort" + "strings" + "sync" + "time" + + "golang.org/x/oauth2" + "golang.org/x/oauth2/google" + + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +// gcpScope is the OAuth2 scope minted for Vertex AI service-account auth. +const gcpScope = "https://www.googleapis.com/auth/cloud-platform" + +// gcpTokenTimeout bounds each GCP token mint/refresh HTTP call so a slow or +// unreachable token endpoint can't block the request indefinitely. +const gcpTokenTimeout = 10 * time.Second + +// ID is the registry key for this middleware. +const ID = "llm_router" + +// Version is reported via Middleware.Version(). +const Version = "1.0.0" + +const ( + denyCodeNotRoutable = "llm_policy.model_not_routable" + denyReasonNotRoutable = "model_not_routable" + denyCodeNoAuthorisedRoute = "llm_policy.no_authorised_provider" + denyReasonNoAuthorisedRoute = "no_authorised_provider" + //nolint:gosec // deny code label, not a credential + denyCodeUpstreamAuth = "llm_policy.upstream_auth_failed" + denyCodeUnmeterable = "llm_policy.unmeterable_publisher" + denyReasonUnmeterable = "unmeterable_publisher" +) + +// strippedAuthHeaders is the closed list of vendor authentication +// credentials the router clears before injecting the provider-specific +// credential. Strictly auth headers — vendor-specific metadata +// (anthropic-version, openai-organization, openai-project, etc.) is +// NOT stripped because the client SDK sets those and the upstream +// requires them (e.g. Anthropic returns 400 without +// anthropic-version). Each entry is canonicalised by Go's +// http.Header.Del/Set, so listing the canonical shapes here is +// sufficient. +var strippedAuthHeaders = []string{ + "Authorization", // OpenAI, OpenAI-compatible, most vendors, Bedrock bearer + "Proxy-Authorization", // upstream proxy auth (defense-in-depth) + "x-api-key", // Anthropic + "api-key", // Azure OpenAI + "X-Amz-Date", // AWS SigV4 — strip client-supplied AWS signing material + "X-Amz-Security-Token", + "X-Amz-Content-Sha256", +} + +// Middleware routes requests to upstream LLM providers based on the +// llm.model metadata emitted by llm_request_parser. +type Middleware struct { + cfg Config + // tokenSrc caches one auto-refreshing OAuth2 TokenSource per GCP + // service-account key (keyed by a hash of the key material), so Vertex + // token minting happens once and refreshes are amortised across requests. + tokenMu sync.Mutex + tokenSrc map[string]oauth2.TokenSource +} + +// New constructs a Middleware with the supplied configuration. Empty +// or nil Providers slice yields a router that denies every request as +// not-routable. +func New(cfg Config) *Middleware { + return &Middleware{cfg: cfg, tokenSrc: map[string]oauth2.TokenSource{}} +} + +// 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 router only consults +// the metadata emitted by llm_request_parser. +func (m *Middleware) AcceptedContentTypes() []string { return nil } + +// MetadataKeys is the closed set of metadata keys this middleware may +// emit. The accumulator drops anything outside this allowlist. +func (m *Middleware) MetadataKeys() []string { + return []string{ + middleware.KeyLLMResolvedProviderID, + middleware.KeyLLMAuthorisingGroups, + middleware.KeyLLMPolicyDecision, + middleware.KeyLLMPolicyReason, + } +} + +// MutationsSupported reports that the middleware emits header and +// upstream-rewrite mutations. +func (m *Middleware) MutationsSupported() bool { return true } + +// Close releases resources owned by the middleware. The router is +// stateless, so this is a no-op. +func (m *Middleware) Close() error { return nil } + +// matchOutcome captures why matchRoute returned what it did so the +// caller can distinguish "no provider knows this model" from "providers +// know it but none authorise this peer's groups". +type matchOutcome int + +const ( + matchOutcomeFound matchOutcome = iota + matchOutcomeUnknownModel + matchOutcomeUnauthorised +) + +// Invoke resolves the model to a provider authorised for the caller's +// groups, strips known vendor auth headers, and injects the route's +// auth header. Unknown models deny with model_not_routable; models +// known to a provider that no policy authorises for the caller deny +// with no_authorised_provider. +func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) { + // Vertex AI carries the model in the URL path, not the body, and is + // selected by path rather than by the model/vendor table. Route it before + // the model lookup so a model the parser extracted from the path can't be + // claimed by a same-vendor direct provider (e.g. claude-* on api.anthropic.com). + reqPath := requestPath(in.URL) + if isVertexPath(reqPath) { + model, _ := lookupMetadata(in.Metadata, middleware.KeyLLMModel) + // The request parser emits no llm.provider for a Vertex publisher it + // can't parse (e.g. google/gemini). Forwarding such a request would + // bypass token/budget metering, so deny it rather than serve it + // unmetered. + if vendor, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider); vendor == "" { + return denyUnmeterable(), nil + } + route, outcome := m.matchVertex(reqPath, model, in.UserGroups) + switch outcome { + case matchOutcomeFound: + return m.allowWithRoute(route, in.UserGroups), nil + case matchOutcomeUnauthorised: + return denyNoAuthorisedRoute(model), nil + default: + return denyUnknownModel(model), nil + } + } + + // Bedrock likewise carries the model in the URL path (/model/{id}/{action}), + // optionally behind a "/bedrock" gateway-namespace prefix. Route it by path + // before the model lookup; when the prefix is present, strip it from the + // forwarded path so the real Bedrock endpoint receives its native path. + if isBedrockPath(reqPath) { + model, _ := lookupMetadata(in.Metadata, middleware.KeyLLMModel) + native, hadPrefix := splitBedrockNamespace(reqPath) + route, outcome := m.matchBedrock(native, model, in.UserGroups) + switch outcome { + case matchOutcomeFound: + out := m.allowWithRoute(route, in.UserGroups) + if hadPrefix && out.Mutations != nil && out.Mutations.RewriteUpstream != nil { + out.Mutations.RewriteUpstream.StripPathPrefix = bedrockNamespacePrefix + } + return out, nil + case matchOutcomeUnauthorised: + return denyNoAuthorisedRoute(model), nil + default: + return denyUnknownModel(model), nil + } + } + + model, ok := lookupMetadata(in.Metadata, middleware.KeyLLMModel) + if !ok || model == "" { + // Non-inference endpoints (model listing) carry no model but still + // need rewriting from the synth placeholder to a real upstream; + // clients such as Codex call GET /v1/models at startup to enumerate + // availability and read a 403 as "model unavailable". + route, outcome := m.matchModelless(requestPath(in.URL), in.UserGroups) + switch outcome { + case matchOutcomeFound: + return m.allowWithRoute(route, in.UserGroups), nil + case matchOutcomeUnauthorised: + // A recognised model-less endpoint exists but no provider + // authorises the caller — deny as an authorisation failure + // rather than masking it as a missing model. + return denyNoAuthorisedRoute(model), nil + default: + return denyMissingModel(), nil + } + } + + vendor, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider) + route, outcome := m.matchRoute(model, vendor, requestPath(in.URL), in.UserGroups) + switch outcome { + case matchOutcomeFound: + return m.allowWithRoute(route, in.UserGroups), nil + case matchOutcomeUnauthorised: + return denyNoAuthorisedRoute(model), nil + default: + return denyUnknownModel(model), nil + } +} + +// matchRoute returns the ProviderRoute that should serve the given +// model + request path for a caller in the given user-groups. Selection +// is: +// +// 1. Filter the configured providers to those whose Models list +// contains the model. +// 2. Filter the model-matched candidates to those whose +// AllowedGroupIDs intersect the caller's UserGroups. A route with +// no AllowedGroupIDs is the catch-all: it stays in the list. If +// the model was known but no candidate is authorised for this +// peer, return matchOutcomeUnauthorised so the caller can emit +// the dedicated no_authorised_provider deny code. +// 3. Vendor precedence: when the request carries a detected vendor +// (llm.provider) and at least one candidate is the same vendor, +// drop the rest — a vendor-tagged request must never cross to +// another vendor's route (e.g. an Anthropic call landing on an +// OpenAI-compatible gateway that also claims the model). +// 4. Model precedence over path: a route that explicitly lists the +// model beats a catch-all (empty Models) gateway. +// 5. Disambiguate the survivors by URL path prefix: longest +// UpstreamPath that prefix-matches the request path wins; an empty +// UpstreamPath is the catchall. If none prefix-matches, fall back +// to declaration order so the model stays routable. +func (m *Middleware) matchRoute(model, vendor, reqPath string, userGroups []string) (ProviderRoute, matchOutcome) { + var modelMatched []ProviderRoute + for _, route := range m.cfg.Providers { + if routeClaimsModel(route, model) { + modelMatched = append(modelMatched, route) + } + } + if len(modelMatched) == 0 { + return ProviderRoute{}, matchOutcomeUnknownModel + } + + // Vendor pinning runs BEFORE the group filter so a request the parser + // tagged with a vendor can never cross to another vendor's route — not + // even an authorised one. Narrow to same-vendor routes when any + // model-matched route declares that vendor; setups with no vendor tag on + // any route fall through unchanged. After narrowing, if no same-vendor + // route authorises the caller, that's matchOutcomeUnauthorised (no + // cross-vendor fallback). + if vendor != "" { + if vendorMatched := matchingVendor(modelMatched, vendor); len(vendorMatched) > 0 { + modelMatched = vendorMatched + } + } + + var candidates []ProviderRoute + for _, route := range modelMatched { + if routeAuthorisesGroups(route, userGroups) { + candidates = append(candidates, route) + } + } + if len(candidates) == 0 { + return ProviderRoute{}, matchOutcomeUnauthorised + } + + // Model routing takes precedence over path. A route that explicitly + // lists the model must beat a catch-all (empty Models) gateway that + // claims every model — otherwise an Anthropic request can fall through + // to an OpenAI-compatible gateway declared earlier. Only when no + // candidate explicitly claims the model do the catch-alls compete, and + // the path-prefix tiebreak applies within whichever tier wins. + if explicit := explicitlyClaiming(candidates, model); len(explicit) > 0 { + candidates = explicit + } + if len(candidates) == 1 { + return candidates[0], matchOutcomeFound + } + + best := candidates[0] + bestLen := -1 + for _, c := range candidates { + if !pathPrefixMatches(c.UpstreamPath, reqPath) { + continue + } + if len(c.UpstreamPath) > bestLen { + best = c + bestLen = len(c.UpstreamPath) + } + } + return best, matchOutcomeFound +} + +// isModelLessPath reports whether reqPath is a known OpenAI-shaped +// non-inference endpoint that legitimately carries no model in its +// request (the model-listing endpoints). These must route to an upstream +// rather than deny, so model enumeration works end to end. +func isModelLessPath(reqPath string) bool { + return reqPath == "/v1/models" || strings.HasPrefix(reqPath, "/v1/models/") +} + +// isVertexPath reports whether reqPath is a Google Vertex AI publisher +// endpoint: /v1/projects/{project}/locations/{region}/publishers/{publisher}/ +// models/{model}:{action}. The model + vendor live in the path, so these +// requests are routed by path to the Vertex provider rather than by model. +func isVertexPath(reqPath string) bool { + return strings.HasPrefix(reqPath, "/v1/projects/") && + strings.Contains(reqPath, "/publishers/") && + strings.Contains(reqPath, "/models/") +} + +// bedrockNamespacePrefix is an optional gateway-namespace prefix some clients +// place before the native Bedrock path to disambiguate it from other providers +// that also use "/model/...". It is stripped before forwarding upstream. +const bedrockNamespacePrefix = "/bedrock" + +// splitBedrockNamespace removes an optional "/bedrock" namespace prefix, +// returning the native Bedrock path and whether the prefix was present. +func splitBedrockNamespace(reqPath string) (string, bool) { + if strings.HasPrefix(reqPath, bedrockNamespacePrefix+"/") { + return strings.TrimPrefix(reqPath, bedrockNamespacePrefix), true + } + return reqPath, false +} + +// isBedrockPath reports whether reqPath is an AWS Bedrock runtime model +// endpoint: /model/{modelId}/{action} where action is invoke, +// invoke-with-response-stream, converse, or converse-stream — optionally behind +// a "/bedrock" gateway-namespace prefix. The model lives in the path, so these +// requests are routed by path to the Bedrock provider. +func isBedrockPath(reqPath string) bool { + native, _ := splitBedrockNamespace(reqPath) + if !strings.HasPrefix(native, "/model/") { + return false + } + return strings.HasSuffix(native, "/invoke") || + strings.HasSuffix(native, "/invoke-with-response-stream") || + strings.HasSuffix(native, "/converse") || + strings.HasSuffix(native, "/converse-stream") +} + +// matchVertex selects the Vertex provider authorised for the caller's groups +// and claiming the requested model. +func (m *Middleware) matchVertex(reqPath, model string, userGroups []string) (ProviderRoute, matchOutcome) { + return m.matchPathRoute(reqPath, model, userGroups, func(r ProviderRoute) bool { return r.Vertex }) +} + +// matchBedrock selects the Bedrock provider authorised for the caller's groups +// and claiming the requested model. +func (m *Middleware) matchBedrock(reqPath, model string, userGroups []string) (ProviderRoute, matchOutcome) { + return m.matchPathRoute(reqPath, model, userGroups, func(r ProviderRoute) bool { return r.Bedrock }) +} + +// matchPathRoute selects a path-routed provider (Vertex/Bedrock). These carry +// the model in the URL, so the model/vendor table is bypassed — but the route's +// configured Models allowlist is still enforced (empty Models = catch-all) so a +// provider credential can't be used for models the operator didn't authorise. +// Returns matchOutcomeUnauthorised when no style route authorises the caller's +// groups, matchOutcomeUnknownModel when an authorised route exists but none +// claims the model (or no style route exists at all), else the chosen route +// (longest UpstreamPath prefix-match wins among multiple). +func (m *Middleware) matchPathRoute(reqPath, model string, userGroups []string, isStyle func(ProviderRoute) bool) (ProviderRoute, matchOutcome) { + var styled []ProviderRoute + for _, route := range m.cfg.Providers { + if isStyle(route) { + styled = append(styled, route) + } + } + if len(styled) == 0 { + return ProviderRoute{}, matchOutcomeUnknownModel + } + + var authorised []ProviderRoute + for _, route := range styled { + if routeAuthorisesGroups(route, userGroups) { + authorised = append(authorised, route) + } + } + if len(authorised) == 0 { + return ProviderRoute{}, matchOutcomeUnauthorised + } + + var candidates []ProviderRoute + for _, route := range authorised { + if routeClaimsModel(route, model) { + candidates = append(candidates, route) + } + } + if len(candidates) == 0 { + return ProviderRoute{}, matchOutcomeUnknownModel + } + if len(candidates) == 1 { + return candidates[0], matchOutcomeFound + } + + best := candidates[0] + bestLen := -1 + for _, c := range candidates { + if !pathPrefixMatches(c.UpstreamPath, reqPath) { + continue + } + if len(c.UpstreamPath) > bestLen { + best = c + bestLen = len(c.UpstreamPath) + } + } + return best, matchOutcomeFound +} + +// matchModelless selects a route for a non-inference, model-less request. +// It mirrors matchRoute's group-authorisation filter and path-prefix +// tiebreak but skips the per-model filter, since any provider the caller's +// groups authorise can serve a model-listing request. Returns +// matchOutcomeFound with the chosen route (single authorised provider wins +// outright; multiple fall to the longest UpstreamPath prefix-match, then +// declaration order), matchOutcomeUnauthorised when no provider authorises +// the caller, or matchOutcomeUnknownModel when the path isn't a recognised +// model-less endpoint. +func (m *Middleware) matchModelless(reqPath string, userGroups []string) (ProviderRoute, matchOutcome) { + if !isModelLessPath(reqPath) { + return ProviderRoute{}, matchOutcomeUnknownModel + } + var candidates []ProviderRoute + for _, route := range m.cfg.Providers { + // Vertex/Bedrock are path-routed and don't serve OpenAI-style + // model-listing endpoints; including them here could rewrite a + // GET /v1/models to an upstream that 404s it. + if route.Vertex || route.Bedrock { + continue + } + if routeAuthorisesGroups(route, userGroups) { + candidates = append(candidates, route) + } + } + if len(candidates) == 0 { + return ProviderRoute{}, matchOutcomeUnauthorised + } + if len(candidates) == 1 { + return candidates[0], matchOutcomeFound + } + + best := candidates[0] + bestLen := -1 + for _, c := range candidates { + if !pathPrefixMatches(c.UpstreamPath, reqPath) { + continue + } + if len(c.UpstreamPath) > bestLen { + best = c + bestLen = len(c.UpstreamPath) + } + } + return best, matchOutcomeFound +} + +// routeAuthorisesGroups reports whether the route's AllowedGroupIDs +// intersect the caller's userGroups. A route with empty AllowedGroupIDs +// is unreachable: the synthesiser only emits routes bound to at least +// one enabled policy, so an empty list signals a misconfiguration that +// must not be allowed to fall through. +func routeAuthorisesGroups(r ProviderRoute, userGroups []string) bool { + for _, ug := range userGroups { + for _, ag := range r.AllowedGroupIDs { + if ug == ag { + return true + } + } + } + return false +} + +// authorisingGroupsCSV returns the sorted, deduplicated comma-separated +// intersection of routeGroups and userGroups — i.e. the groups that +// actually authorise the resolved route for this caller. Returns the +// empty string when the intersection is empty (shouldn't happen on the +// allow path, but defensive). +func authorisingGroupsCSV(routeGroups, userGroups []string) string { + if len(routeGroups) == 0 || len(userGroups) == 0 { + return "" + } + allowed := make(map[string]struct{}, len(routeGroups)) + for _, g := range routeGroups { + allowed[g] = struct{}{} + } + seen := make(map[string]struct{}, len(userGroups)) + out := make([]string, 0, len(userGroups)) + for _, ug := range userGroups { + if _, ok := allowed[ug]; !ok { + continue + } + if _, dup := seen[ug]; dup { + continue + } + seen[ug] = struct{}{} + out = append(out, ug) + } + if len(out) == 0 { + return "" + } + sort.Strings(out) + return strings.Join(out, ",") +} + +// matchingVendor returns the subset of routes whose Vendor equals the +// request's detected vendor. Routes with an empty Vendor never match — an +// untagged route can't be asserted to speak the request's surface, so it +// stays out of the vendor-filtered set (but remains eligible via the +// fall-through when no route matches the vendor at all). +func matchingVendor(routes []ProviderRoute, vendor string) []ProviderRoute { + var out []ProviderRoute + for _, r := range routes { + if r.Vendor == vendor { + out = append(out, r) + } + } + return out +} + +// explicitlyClaiming returns the subset of routes whose Models list +// names the model exactly. Catch-all routes (empty Models) are excluded, +// so callers can prefer a provider that genuinely declares the model over +// a gateway that claims everything. +func explicitlyClaiming(routes []ProviderRoute, model string) []ProviderRoute { + var out []ProviderRoute + for _, r := range routes { + for _, candidate := range r.Models { + if candidate == model { + out = append(out, r) + break + } + } + } + return out +} + +// routeClaimsModel reports whether the route's Models list contains +// the given model identifier. An empty Models list is treated as +// "claim every model" — used by gateway-style providers (LiteLLM, +// custom OpenAI-compatible endpoints) that proxy an open-ended set of +// upstream models the operator can't enumerate in NetBird's provider +// config. +func routeClaimsModel(route ProviderRoute, model string) bool { + if len(route.Models) == 0 { + return true + } + for _, candidate := range route.Models { + if candidate == model { + return true + } + } + return false +} + +// pathPrefixMatches reports whether upstreamPath matches reqPath on a path- +// segment boundary: an exact match, or reqPath continuing after +// upstreamPath at a "/" separator. This avoids a sibling base like +// "/openai" spuriously matching "/openai-test". An empty (or "/") +// upstreamPath always matches (catchall). +func pathPrefixMatches(upstreamPath, reqPath string) bool { + if upstreamPath == "" || upstreamPath == "/" { + return true + } + upstreamPath = strings.TrimRight(upstreamPath, "/") + return reqPath == upstreamPath || strings.HasPrefix(reqPath, upstreamPath+"/") +} + +// requestPath extracts the path component from an Input.URL string +// (which is r.URL.String() — typically "/path?query"). Returns the +// raw input on parse failure so the prefix check can still operate on +// the unparsed value. +func requestPath(raw string) string { + if raw == "" { + return "" + } + parsed, err := url.Parse(raw) + if err != nil { + return raw + } + return parsed.Path +} + +// allowWithRoute builds the Output for a successful route match. The +// returned Mutations carry the upstream rewrite plus — riding on it — +// the StripHeaders list and the AuthHeader to inject. +// +// The strip + inject MUST go through UpstreamRewrite (not HeadersAdd / +// HeadersRemove) because the framework's mutation gate runs every +// header change through a denylist that blocks Authorization, +// Cookie, etc. — exactly the headers the router is replacing. The +// proxy's upstream-build path applies AuthHeader / StripHeaders +// directly, bypassing the denylist by virtue of being a trusted +// proxy operation rather than an arbitrary middleware mutation. +// +// Emits the authorising-groups intersection alongside the resolved +// provider id so identity-stamping middlewares (llm_identity_inject) +// tag the request with ONLY the groups that authorised this specific +// route — not every group the peer happens to be in. +func (m *Middleware) allowWithRoute(route ProviderRoute, userGroups []string) *middleware.Output { + rewrite := &middleware.UpstreamRewrite{ + Scheme: route.UpstreamScheme, + Host: route.UpstreamHost, + // UpstreamPath is the path component the operator pasted on + // the provider record (e.g. "/v1/{account}/{gateway}/compat" + // for Cloudflare AI Gateway). Carrying it on the rewrite so + // the proxy's URL composer joins it with the agent's request + // path — without this, the operator's configured upstream + // path is silently dropped and the gateway returns a 4xx for + // the malformed URL. Empty value leaves the original + // target's path untouched. + Path: route.UpstreamPath, + StripHeaders: append([]string(nil), strippedAuthHeaders...), + } + authValue := route.AuthHeaderValue + if route.GCPServiceAccountKeyB64 != "" { + // Mint a short-lived OAuth2 token from the service-account key at + // request time (cached + auto-refreshed) instead of a static value. + bearer, err := m.gcpBearer(route.GCPServiceAccountKeyB64) + if err != nil { + return denyUpstreamAuth() + } + authValue = bearer + } + if route.AuthHeaderName != "" && authValue != "" { + rewrite.AuthHeader = &middleware.AuthHeader{ + Name: route.AuthHeaderName, + Value: authValue, + } + } + return &middleware.Output{ + Decision: middleware.DecisionAllow, + Mutations: &middleware.Mutations{RewriteUpstream: rewrite}, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMResolvedProviderID, Value: route.ID}, + {Key: middleware.KeyLLMAuthorisingGroups, Value: authorisingGroupsCSV(route.AllowedGroupIDs, userGroups)}, + {Key: middleware.KeyLLMPolicyDecision, Value: "allow"}, + }, + } +} + +// gcpBearer returns a "Bearer " value minted from a base64-encoded GCP +// service-account key, using a cached, auto-refreshing token source. +func (m *Middleware) gcpBearer(saKeyB64 string) (string, error) { + ts, err := m.gcpTokenSource(saKeyB64) + if err != nil { + return "", err + } + tok, err := ts.Token() + if err != nil { + return "", fmt.Errorf("mint gcp token: %w", err) + } + return "Bearer " + tok.AccessToken, nil +} + +// gcpTokenSource returns the cached TokenSource for the given service-account +// key, building it (decode base64 → parse JSON → cloud-platform scope) on first +// use. The returned source caches the token and refreshes it before expiry. +func (m *Middleware) gcpTokenSource(saKeyB64 string) (oauth2.TokenSource, error) { + sum := sha256.Sum256([]byte(saKeyB64)) + key := hex.EncodeToString(sum[:]) + + m.tokenMu.Lock() + defer m.tokenMu.Unlock() + if m.tokenSrc == nil { + m.tokenSrc = map[string]oauth2.TokenSource{} + } + if ts, ok := m.tokenSrc[key]; ok { + return ts, nil + } + jsonKey, err := base64.StdEncoding.DecodeString(strings.TrimSpace(saKeyB64)) + if err != nil { + return nil, fmt.Errorf("decode gcp service-account key: %w", err) + } + conf, err := google.JWTConfigFromJSON(jsonKey, gcpScope) + if err != nil { + return nil, fmt.Errorf("parse gcp service-account key: %w", err) + } + // Bound mint/refresh with a timeout HTTP client so a slow token endpoint + // can't hang the request. The oauth2 library uses this client for the + // lifetime of the (auto-refreshing) source. + ctx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{Timeout: gcpTokenTimeout}) + ts := conf.TokenSource(ctx) + m.tokenSrc[key] = ts + return ts, nil +} + +// denyUpstreamAuth is returned when the router cannot obtain the upstream +// credential (e.g. a malformed service-account key or an unreachable token +// endpoint). It surfaces as a 502 — an upstream problem, not a policy denial. +func denyUpstreamAuth() *middleware.Output { + return &middleware.Output{ + Decision: middleware.DecisionDeny, + DenyStatus: 502, + DenyReason: &middleware.DenyReason{ + Code: denyCodeUpstreamAuth, + Message: "could not obtain upstream credential", + }, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMPolicyDecision, Value: "deny"}, + {Key: middleware.KeyLLMPolicyReason, Value: "upstream_auth_failed"}, + }, + } +} + +// denyUnmeterable returns the deny envelope for a path-routed request whose +// publisher has no parser surface, so its usage can't be metered. Serving it +// would bypass token/budget caps, so it is rejected with a 403. +func denyUnmeterable() *middleware.Output { + return &middleware.Output{ + Decision: middleware.DecisionDeny, + DenyStatus: 403, + DenyReason: &middleware.DenyReason{ + Code: denyCodeUnmeterable, + Message: "request publisher is not supported for metering", + }, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMPolicyDecision, Value: "deny"}, + {Key: middleware.KeyLLMPolicyReason, Value: denyReasonUnmeterable}, + }, + } +} + +// denyMissingModel returns the deny envelope for a request whose +// envelope has no llm.model metadata. +func denyMissingModel() *middleware.Output { + return &middleware.Output{ + Decision: middleware.DecisionDeny, + DenyStatus: 403, + DenyReason: &middleware.DenyReason{ + Code: denyCodeNotRoutable, + Message: "missing llm.model on request envelope", + }, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMPolicyDecision, Value: "deny"}, + {Key: middleware.KeyLLMPolicyReason, Value: denyReasonNotRoutable}, + }, + } +} + +// denyUnknownModel returns the deny envelope for a model that no +// configured provider claims. +func denyUnknownModel(model string) *middleware.Output { + return &middleware.Output{ + Decision: middleware.DecisionDeny, + DenyStatus: 403, + DenyReason: &middleware.DenyReason{ + Code: denyCodeNotRoutable, + Message: fmt.Sprintf("no provider configured for model %s", model), + Details: map[string]string{"model": model}, + }, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMPolicyDecision, Value: "deny"}, + {Key: middleware.KeyLLMPolicyReason, Value: denyReasonNotRoutable}, + }, + } +} + +// denyNoAuthorisedRoute returns the deny envelope for a model that one +// or more providers claim, but where no policy authorises the caller's +// groups for any of those providers. +func denyNoAuthorisedRoute(model string) *middleware.Output { + return &middleware.Output{ + Decision: middleware.DecisionDeny, + DenyStatus: 403, + DenyReason: &middleware.DenyReason{ + Code: denyCodeNoAuthorisedRoute, + Message: fmt.Sprintf("no policy authorises model %s for the caller's groups", model), + Details: map[string]string{"model": model}, + }, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMPolicyDecision, Value: "deny"}, + {Key: middleware.KeyLLMPolicyReason, Value: denyReasonNoAuthorisedRoute}, + }, + } +} + +// lookupMetadata returns the value for key plus a presence flag so +// callers can distinguish absent from empty. +func lookupMetadata(meta []middleware.KV, key string) (string, bool) { + for _, kv := range meta { + if kv.Key == key { + return kv.Value, true + } + } + return "", false +} diff --git a/proxy/internal/middleware/builtin/llm_router/middleware_test.go b/proxy/internal/middleware/builtin/llm_router/middleware_test.go new file mode 100644 index 000000000..8ae03c5ba --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_router/middleware_test.go @@ -0,0 +1,840 @@ +package llm_router + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +// metaValue returns the value for the first KV with the given key. +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 +} + +// defaultTestGroup is the group id used by routes and inputs in tests +// that don't specifically exercise the group-filter logic. Pairing it +// with the same id on every test route keeps the legacy assertions +// focused on routing/path behaviour without each one having to bake in +// its own ACL. +const defaultTestGroup = "grp-test" + +// newInputWithModel returns an Input carrying llm.model in its metadata +// bag, mimicking the post-llm_request_parser state the router observes +// in production. UserGroups is populated with defaultTestGroup so the +// router's group-filter pass authorises any test route whose +// AllowedGroupIDs contains the same id. +func newInputWithModel(model string) *middleware.Input { + return &middleware.Input{ + Slot: middleware.SlotOnRequest, + Metadata: []middleware.KV{{Key: middleware.KeyLLMModel, Value: model}}, + UserGroups: []string{defaultTestGroup}, + } +} + +// newInputWithModelAndURL returns an Input carrying both llm.model and +// a request URL so router tests can exercise path-based disambiguation. +func newInputWithModelAndURL(model, reqURL string) *middleware.Input { + in := newInputWithModel(model) + in.URL = reqURL + return in +} + +func TestMiddlewareIdentity(t *testing.T) { + mw := New(Config{}) + assert.Equal(t, ID, mw.ID(), "middleware ID must be llm_router") + assert.Equal(t, Version, mw.Version(), "version must match the constant") + assert.Equal(t, middleware.SlotOnRequest, mw.Slot(), "router must run in SlotOnRequest") + assert.True(t, mw.MutationsSupported(), "router must declare mutations support") + assert.Nil(t, mw.AcceptedContentTypes(), "router does not inspect bodies") + assert.ElementsMatch(t, + []string{ + middleware.KeyLLMResolvedProviderID, + middleware.KeyLLMAuthorisingGroups, + middleware.KeyLLMPolicyDecision, + middleware.KeyLLMPolicyReason, + }, + mw.MetadataKeys(), + "metadata key allowlist must match the spec", + ) + require.NoError(t, mw.Close()) +} + +func TestRouter_HappyPath(t *testing.T) { + route := ProviderRoute{ + ID: "openai-prod", + Models: []string{"gpt-4o"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer sk-test-123", + } + mw := New(Config{Providers: []ProviderRoute{route}}) + + out, err := mw.Invoke(context.Background(), newInputWithModel("gpt-4o")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "matched model must allow") + + require.NotNil(t, out.Mutations, "matched route must emit mutations") + rewrite := out.Mutations.RewriteUpstream + require.NotNil(t, rewrite, "matched route must emit upstream rewrite") + assert.Equal(t, "https", rewrite.Scheme, "rewrite scheme must come from the matched route") + assert.Equal(t, "api.openai.com", rewrite.Host, "rewrite host must come from the matched route") + + assert.ElementsMatch(t, strippedAuthHeaders, rewrite.StripHeaders, + "strip list rides on UpstreamRewrite (bypasses framework denylist) and must cover every known vendor auth header") + require.NotNil(t, rewrite.AuthHeader, "router must inject the auth header via the rewrite (not HeadersAdd) so the proxy bypasses the denylist") + assert.Equal(t, "Authorization", rewrite.AuthHeader.Name, "injected header name must come from the route") + assert.Equal(t, "Bearer sk-test-123", rewrite.AuthHeader.Value, "injected header value must come from the route") + assert.Empty(t, out.Mutations.HeadersAdd, "router must not use HeadersAdd; auth flows through UpstreamRewrite.AuthHeader") + assert.Empty(t, out.Mutations.HeadersRemove, "router must not use HeadersRemove; strip flows through UpstreamRewrite.StripHeaders") + + resolved, ok := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID) + require.True(t, ok, "router must emit llm.resolved_provider_id on a match") + assert.Equal(t, "openai-prod", resolved, "resolved provider id must be the matched route's ID") + dec, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision) + assert.Equal(t, "allow", dec, "decision metadata must be allow on a match") +} + +func TestRouter_MissingModel(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "openai-prod", + Models: []string{"gpt-4o"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + }}}) + + out, err := mw.Invoke(context.Background(), &middleware.Input{Slot: middleware.SlotOnRequest}) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionDeny, out.Decision, "missing llm.model must deny") + assert.Equal(t, 403, out.DenyStatus, "deny status must be 403") + require.NotNil(t, out.DenyReason, "deny reason must be populated") + assert.Equal(t, "llm_policy.model_not_routable", out.DenyReason.Code, "deny code must be model_not_routable") + assert.Equal(t, "missing llm.model on request envelope", out.DenyReason.Message, "deny message must match spec") + + dec, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision) + assert.Equal(t, "deny", dec, "decision metadata must be deny") + reason, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyReason) + assert.Equal(t, "model_not_routable", reason, "reason metadata must be model_not_routable") +} + +// newModellessInput returns an Input with no llm.model and the given +// request path, mimicking a GET /v1/models call (which carries no body +// from which a model could be parsed). UserGroups matches defaultTestGroup. +func newModellessInput(reqURL string) *middleware.Input { + return &middleware.Input{ + Slot: middleware.SlotOnRequest, + URL: reqURL, + UserGroups: []string{defaultTestGroup}, + } +} + +func TestRouter_ModelLessPath_RoutesToAuthorisedProvider(t *testing.T) { + route := ProviderRoute{ + ID: "openai-prod", + Models: []string{"gpt-4o"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + } + mw := New(Config{Providers: []ProviderRoute{route}}) + + out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models?client_version=1")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "GET /v1/models must pass through, not deny") + require.NotNil(t, out.Mutations, "a pass-through must rewrite the upstream") + require.NotNil(t, out.Mutations.RewriteUpstream, "model-less route must still rewrite to the real upstream") + assert.Equal(t, "api.openai.com", out.Mutations.RewriteUpstream.Host, "must target the authorised provider's host") + + provider, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID) + assert.Equal(t, "openai-prod", provider, "resolved provider must be the authorised route") +} + +func TestRouter_ModelLessPath_MultiProviderDeclarationOrder(t *testing.T) { + first := ProviderRoute{ + ID: "openai-a", + Models: []string{"gpt-4o"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "a.example.com", + } + second := ProviderRoute{ + ID: "openai-b", + Models: []string{"gpt-4o-mini"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "b.example.com", + } + mw := New(Config{Providers: []ProviderRoute{first, second}}) + + out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "model-less path must pass through with multiple providers") + provider, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID) + assert.Equal(t, "openai-a", provider, "no path-prefix match falls back to declaration order") +} + +func TestRouter_ModelLessPath_UnauthorisedDenies(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "openai-prod", + Models: []string{"gpt-4o"}, + AllowedGroupIDs: []string{"some-other-group"}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + }}}) + + out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionDeny, out.Decision, "no provider authorising the caller must still deny") +} + +func TestRouter_NonModelLessBodilessStillDenies(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "openai-prod", + Models: []string{"gpt-4o"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + }}}) + + // A bodiless POST to an inference path has no model and is NOT a + // model-less endpoint, so it must keep denying. + out, err := mw.Invoke(context.Background(), newModellessInput("/v1/responses")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionDeny, out.Decision, "bodiless inference request must still deny") + assert.Equal(t, "llm_policy.model_not_routable", out.DenyReason.Code, "deny code stays model_not_routable") +} + +// TestRouter_ExplicitModelBeatsCatchallGateway is the regression guard +// for multi-provider misrouting: a catch-all (empty Models) OpenAI-compat +// gateway declared first must NOT swallow a model an explicit provider +// claims. Anthropic's claude request must reach the Anthropic route even +// though the gateway claims every model and wins declaration order. +func TestRouter_ExplicitModelBeatsCatchallGateway(t *testing.T) { + gateway := ProviderRoute{ + ID: "openai-gateway", + Models: nil, // catch-all: claims every model + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + } + anthropic := ProviderRoute{ + ID: "anthropic-prod", + Models: []string{"claude-opus-4"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.anthropic.com", + } + // Gateway declared first to prove explicit claim beats declaration order. + mw := New(Config{Providers: []ProviderRoute{gateway, anthropic}}) + + out, err := mw.Invoke(context.Background(), newInputWithModelAndURL("claude-opus-4", "/v1/messages")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "explicit-model request must route, not deny") + require.NotNil(t, out.Mutations) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "api.anthropic.com", out.Mutations.RewriteUpstream.Host, "claude must reach the explicit Anthropic route, not the catch-all gateway") + + provider, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID) + assert.Equal(t, "anthropic-prod", provider, "resolved provider must be the explicit Anthropic route") +} + +// TestRouter_CatchallStillServesUnlistedModel confirms the catch-all +// gateway still wins models no explicit provider claims (its whole point). +func TestRouter_CatchallStillServesUnlistedModel(t *testing.T) { + gateway := ProviderRoute{ + ID: "openai-gateway", + Models: nil, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "gateway.example.com", + } + anthropic := ProviderRoute{ + ID: "anthropic-prod", + Models: []string{"claude-opus-4"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.anthropic.com", + } + mw := New(Config{Providers: []ProviderRoute{gateway, anthropic}}) + + out, err := mw.Invoke(context.Background(), newInputWithModelAndURL("some-exotic-model", "/v1/chat/completions")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "unlisted model must still route via the catch-all") + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "gateway.example.com", out.Mutations.RewriteUpstream.Host, "unlisted model falls to the catch-all gateway") +} + +// newInputVendorModelURL returns an Input carrying both the detected +// vendor (llm.provider) and the model, plus a request URL — mimicking the +// post-llm_request_parser state for a real inference call. +func newInputVendorModelURL(vendor, model, reqURL string) *middleware.Input { + return &middleware.Input{ + Slot: middleware.SlotOnRequest, + URL: reqURL, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMProvider, Value: vendor}, + {Key: middleware.KeyLLMModel, Value: model}, + }, + UserGroups: []string{defaultTestGroup}, + } +} + +// TestRouter_VendorKeepsAnthropicOffOpenAIGateway is the regression guard +// for the reported multi-provider break: two catch-all providers (neither +// enumerates models), the OpenAI one declared first. Without vendor +// awareness, a claude request matches both, no path prefixes, and +// declaration order sends it to OpenAI → 502. The detected vendor must +// pin it to the Anthropic route. +func TestRouter_VendorKeepsAnthropicOffOpenAIGateway(t *testing.T) { + openai := ProviderRoute{ + ID: "openai-gw", + Vendor: "openai", + Models: nil, // catch-all + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + } + anthropic := ProviderRoute{ + ID: "anthropic-gw", + Vendor: "anthropic", + Models: nil, // catch-all + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.anthropic.com", + } + mw := New(Config{Providers: []ProviderRoute{openai, anthropic}}) // openai first + + out, err := mw.Invoke(context.Background(), newInputVendorModelURL("anthropic", "claude-opus-4-8", "/v1/messages")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "claude request must route, not deny") + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "api.anthropic.com", out.Mutations.RewriteUpstream.Host, "anthropic vendor must pin to the anthropic route despite openai being declared first") + + provider, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID) + assert.Equal(t, "anthropic-gw", provider) +} + +// TestRouter_VendorKeepsOpenAIOffAnthropic is the reciprocal: an OpenAI +// request must stay on the OpenAI route even when the Anthropic catch-all +// is declared first. +func TestRouter_VendorKeepsOpenAIOffAnthropic(t *testing.T) { + anthropic := ProviderRoute{ + ID: "anthropic-gw", + Vendor: "anthropic", + Models: nil, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.anthropic.com", + } + openai := ProviderRoute{ + ID: "openai-gw", + Vendor: "openai", + Models: nil, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + } + mw := New(Config{Providers: []ProviderRoute{anthropic, openai}}) // anthropic first + + out, err := mw.Invoke(context.Background(), newInputVendorModelURL("openai", "gpt-5.5", "/v1/responses")) + require.NoError(t, err) + require.NotNil(t, out) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "api.openai.com", out.Mutations.RewriteUpstream.Host, "openai vendor must pin to the openai route despite anthropic being declared first") +} + +// TestRouter_VendorAbsentFallsBackToModelPath confirms vendor filtering is +// inert when the request carries no detected vendor: routing then relies on +// model/path as before. +func TestRouter_VendorAbsentFallsBackToModelPath(t *testing.T) { + openai := ProviderRoute{ + ID: "openai-gw", + Vendor: "openai", + Models: []string{"gpt-5.5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + } + mw := New(Config{Providers: []ProviderRoute{openai}}) + + // No llm.provider in metadata — only the model. + out, err := mw.Invoke(context.Background(), newInputWithModel("gpt-5.5")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "explicit-model match must still route with no vendor present") +} + +func TestRouter_UnknownModel(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "openai-prod", + Models: []string{"gpt-4o"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + }}}) + + out, err := mw.Invoke(context.Background(), newInputWithModel("claude-opus-4")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionDeny, out.Decision, "unrouted model must deny") + assert.Equal(t, 403, out.DenyStatus, "deny status must be 403") + require.NotNil(t, out.DenyReason, "deny reason must be populated") + assert.Equal(t, "llm_policy.model_not_routable", out.DenyReason.Code, "deny code must be model_not_routable") + assert.Equal(t, "no provider configured for model claude-opus-4", out.DenyReason.Message, "deny message must reference the offending model") + assert.Equal(t, "claude-opus-4", out.DenyReason.Details["model"], "deny details must include the offending model") +} + +func TestRouter_HeaderStripList(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "openai-prod", + Models: []string{"gpt-4o"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer sk-test-123", + }}}) + + out, err := mw.Invoke(context.Background(), newInputWithModel("gpt-4o")) + require.NoError(t, err) + require.NotNil(t, out) + require.NotNil(t, out.Mutations, "matched route must emit mutations") + + expected := []string{ + "Authorization", + "Proxy-Authorization", + "x-api-key", + "api-key", + } + require.NotNil(t, out.Mutations.RewriteUpstream, "matched route must emit upstream rewrite") + for _, header := range expected { + assert.Contains(t, out.Mutations.RewriteUpstream.StripHeaders, header, + "strip list (on UpstreamRewrite) must include the well-known vendor auth header %s", header) + } + + // Vendor metadata headers MUST NOT be stripped: the client SDK sets them + // and the upstream requires them. Anthropic returns 400 "anthropic-version: + // header is required" if we drop it. Lock the regression. + preserved := []string{"anthropic-version", "openai-organization", "openai-project"} + for _, header := range preserved { + assert.NotContains(t, out.Mutations.RewriteUpstream.StripHeaders, header, + "vendor metadata header %s must NOT be stripped — upstreams require it", header) + } +} + +func TestRouter_FirstMatchWins(t *testing.T) { + first := ProviderRoute{ + ID: "first", + Models: []string{"gpt-4o"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "first.test", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer first", + } + second := ProviderRoute{ + ID: "second", + Models: []string{"gpt-4o"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "second.test", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer second", + } + mw := New(Config{Providers: []ProviderRoute{first, second}}) + + out, err := mw.Invoke(context.Background(), newInputWithModel("gpt-4o")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "duplicate-model match must still allow") + require.NotNil(t, out.Mutations, "matched route must emit mutations") + require.NotNil(t, out.Mutations.RewriteUpstream, "matched route must emit upstream rewrite") + assert.Equal(t, "first.test", out.Mutations.RewriteUpstream.Host, "first-match-wins must pick the earlier route") + + resolved, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID) + assert.Equal(t, "first", resolved, "resolved provider id must be the earlier route's ID") +} + +// TestRouter_PathDisambiguation_PrefixWinsOverCatchall locks in the +// rule the user nailed down: two providers claim the same model, one +// has an UpstreamPath that prefixes the incoming URL, the other has +// no path. The path-prefixed provider wins because the path is a +// strictly more specific match than the empty catchall. +func TestRouter_PathDisambiguation_PrefixWinsOverCatchall(t *testing.T) { + corp := ProviderRoute{ + ID: "corp-openai-compat", + Models: []string{"gpt-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "corp.example.com", + UpstreamPath: "/openai", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer corp", + } + openai := ProviderRoute{ + ID: "openai", + Models: []string{"gpt-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer openai", + } + mw := New(Config{Providers: []ProviderRoute{openai, corp}}) // openai listed first to prove path beats declaration order + + out, err := mw.Invoke(context.Background(), newInputWithModelAndURL("gpt-5", "/openai/v1/chat/completions")) + require.NoError(t, err) + require.NotNil(t, out) + require.Equal(t, middleware.DecisionAllow, out.Decision, "path-prefix match must allow") + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "corp.example.com", out.Mutations.RewriteUpstream.Host, + "path-prefixed provider must beat the catchall when its UpstreamPath is a prefix of the request path") + resolved, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID) + assert.Equal(t, "corp-openai-compat", resolved, "resolved provider id must reflect the path-prefix winner, not the first declared") +} + +// TestRouter_PathDisambiguation_CatchallWhenNoPrefixMatches is the +// inverse: the path-prefixed provider does NOT match the incoming +// path, so the empty-path catchall takes the request. +func TestRouter_PathDisambiguation_CatchallWhenNoPrefixMatches(t *testing.T) { + corp := ProviderRoute{ + ID: "corp-openai-compat", + Models: []string{"gpt-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "corp.example.com", + UpstreamPath: "/openai", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer corp", + } + openai := ProviderRoute{ + ID: "openai", + Models: []string{"gpt-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer openai", + } + mw := New(Config{Providers: []ProviderRoute{corp, openai}}) + + out, err := mw.Invoke(context.Background(), newInputWithModelAndURL("gpt-5", "/v1/chat/completions")) + require.NoError(t, err) + require.NotNil(t, out) + require.Equal(t, middleware.DecisionAllow, out.Decision, "catchall must allow when no path prefix matches") + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "api.openai.com", out.Mutations.RewriteUpstream.Host, + "empty-path catchall must win when the path-prefixed provider's UpstreamPath does not match the request") + resolved, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID) + assert.Equal(t, "openai", resolved, "resolved provider id must be the catchall") +} + +// TestRouter_PathDisambiguation_LongestPrefixWins covers the case +// where multiple providers have non-empty UpstreamPath values that +// both prefix the request — the longer (more specific) one wins. +func TestRouter_PathDisambiguation_LongestPrefixWins(t *testing.T) { + short := ProviderRoute{ + ID: "short-prefix", + Models: []string{"gpt-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "short.example.com", + UpstreamPath: "/openai", + } + long := ProviderRoute{ + ID: "long-prefix", + Models: []string{"gpt-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "long.example.com", + UpstreamPath: "/openai/v1", + } + mw := New(Config{Providers: []ProviderRoute{short, long}}) + + out, err := mw.Invoke(context.Background(), newInputWithModelAndURL("gpt-5", "/openai/v1/chat/completions")) + require.NoError(t, err) + require.NotNil(t, out) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "long.example.com", out.Mutations.RewriteUpstream.Host, + "longest matching UpstreamPath must win — most specific match") +} + +// TestRouter_SingleMatchIgnoresPath proves the path-prefix rule is a +// disambiguation pass, not a gate: when only one provider claims the +// model, it wins regardless of UpstreamPath. Otherwise a path-scoped +// provider would 403 every request whose URL doesn't include the +// path, which would break SDKs configured to hit the gateway root. +func TestRouter_SingleMatchIgnoresPath(t *testing.T) { + only := ProviderRoute{ + ID: "only", + Models: []string{"gpt-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "only.example.com", + UpstreamPath: "/openai", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer only", + } + mw := New(Config{Providers: []ProviderRoute{only}}) + + out, err := mw.Invoke(context.Background(), newInputWithModelAndURL("gpt-5", "/v1/chat/completions")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, + "single model-matching provider must serve the request even when UpstreamPath doesn't prefix the URL — path is a tiebreaker, not a gate") + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "only.example.com", out.Mutations.RewriteUpstream.Host, "the only model-matching provider should be selected") +} + +// TestRouter_PathDisambiguation_FallbackWhenNoPrefixMatches covers +// the multi-candidate edge case where every candidate has a +// non-matching non-empty UpstreamPath. The router falls back to +// declaration order so the model is still routable rather than 403'd. +func TestRouter_PathDisambiguation_FallbackWhenNoPrefixMatches(t *testing.T) { + first := ProviderRoute{ + ID: "first", + Models: []string{"gpt-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "first.example.com", + UpstreamPath: "/openai", + } + second := ProviderRoute{ + ID: "second", + Models: []string{"gpt-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "second.example.com", + UpstreamPath: "/anthropic", + } + mw := New(Config{Providers: []ProviderRoute{first, second}}) + + out, err := mw.Invoke(context.Background(), newInputWithModelAndURL("gpt-5", "/v1/chat/completions")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "no path match among multi-candidates must still allow") + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "first.example.com", out.Mutations.RewriteUpstream.Host, + "when no candidate's UpstreamPath prefix-matches the request, fall back to declaration order") +} + +func TestRouter_FactoryRejectsBadJSON(t *testing.T) { + _, err := Factory{}.New([]byte("{not json")) + require.Error(t, err, "malformed JSON config must be rejected at chain build time") +} + +func TestRouter_FactoryAcceptsEmptyShapes(t *testing.T) { + cases := [][]byte{nil, []byte(""), []byte(" "), []byte("null"), []byte("{}"), []byte("[]")} + for _, raw := range cases { + mw, err := Factory{}.New(raw) + require.NoError(t, err, "empty-shaped config must yield a router with an empty Providers slice") + require.NotNil(t, mw, "factory must return a non-nil middleware on empty config") + + out, invErr := mw.Invoke(context.Background(), newInputWithModel("gpt-4o")) + require.NoError(t, invErr) + assert.Equal(t, middleware.DecisionDeny, out.Decision, + "router with no providers must deny every model as not-routable") + } +} + +// newInputWithModelAndGroups returns an Input carrying llm.model + the +// caller's UserGroups, mimicking the post-auth, post-llm_request_parser +// state the router observes. +func newInputWithModelAndGroups(model string, groups []string) *middleware.Input { + in := newInputWithModel(model) + in.UserGroups = append([]string(nil), groups...) + return in +} + +// TestRouter_GroupFilter_PicksAuthorisedAmongDuplicates pins the Fix A +// behaviour: when two providers claim the same model but each +// authorises a different group, the router must pick the route the +// caller's groups intersect, regardless of declaration order. +func TestRouter_GroupFilter_PicksAuthorisedAmongDuplicates(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{ + { + ID: "openai-marketing", + Models: []string{"gpt-4o-mini"}, + UpstreamScheme: "https", + UpstreamHost: "mkt-openai.example.com", + AllowedGroupIDs: []string{"grp-mkt"}, + }, + { + ID: "openai-engineering", + Models: []string{"gpt-4o-mini"}, + UpstreamScheme: "https", + UpstreamHost: "eng-openai.example.com", + AllowedGroupIDs: []string{"grp-eng"}, + }, + }}) + + out, err := mw.Invoke(context.Background(), + newInputWithModelAndGroups("gpt-4o-mini", []string{"grp-eng"})) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, + "authorised candidate exists; must allow") + + resolved, ok := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID) + require.True(t, ok) + assert.Equal(t, "openai-engineering", resolved, + "router must pick the route whose AllowedGroupIDs intersects the caller's groups, ignoring declaration order") +} + +// TestRouter_GroupFilter_NoIntersection_DeniesNoAuthorisedRoute pins +// the dedicated deny code that fires when the model is known to a +// provider but no candidate is authorised for the caller's groups. +func TestRouter_GroupFilter_NoIntersection_DeniesNoAuthorisedRoute(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "openai-marketing", + Models: []string{"gpt-4o-mini"}, + UpstreamScheme: "https", + UpstreamHost: "mkt-openai.example.com", + AllowedGroupIDs: []string{"grp-mkt"}, + }}}) + + out, err := mw.Invoke(context.Background(), + newInputWithModelAndGroups("gpt-4o-mini", []string{"grp-eng"})) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionDeny, out.Decision, + "model exists but no route authorises grp-eng; must deny") + require.NotNil(t, out.DenyReason) + assert.Equal(t, "llm_policy.no_authorised_provider", out.DenyReason.Code, + "deny code must be no_authorised_provider, not model_not_routable") + assert.Equal(t, "gpt-4o-mini", out.DenyReason.Details["model"], + "deny details must reference the offending model") + + dec, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision) + assert.Equal(t, "deny", dec) + reason, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyReason) + assert.Equal(t, "no_authorised_provider", reason) +} + +// TestRouter_GroupFilter_EmptyAllowedGroupsIsUnreachable pins the +// strict semantics: a route with no AllowedGroupIDs is unreachable. +// The synthesiser only emits policy-bound routes, so an empty ACL +// signals a misconfiguration that must not silently fall through. +func TestRouter_GroupFilter_EmptyAllowedGroupsIsUnreachable(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "openai-shared", + Models: []string{"gpt-4o"}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + // AllowedGroupIDs intentionally left empty. + }}}) + + out, err := mw.Invoke(context.Background(), newInputWithModel("gpt-4o")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionDeny, out.Decision, + "empty AllowedGroupIDs must deny — there is no catch-all for routes without an authorising policy") + require.NotNil(t, out.DenyReason) + assert.Equal(t, "llm_policy.no_authorised_provider", out.DenyReason.Code, + "empty ACL fails the group-filter pass; deny code must reflect that") +} + +// TestRouter_GroupFilter_OverlapTiebreakUnchanged pins that when more +// than one route is authorised for the caller's groups, the existing +// path-prefix tiebreak still decides. Group filtering is a hard gate +// before the tiebreak; it does not change the tiebreak semantics. +func TestRouter_GroupFilter_OverlapTiebreakUnchanged(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{ + { + ID: "openai-a", + Models: []string{"gpt-4o-mini"}, + UpstreamScheme: "https", + UpstreamHost: "a.example.com", + UpstreamPath: "", + AllowedGroupIDs: []string{"grp-eng"}, + }, + { + ID: "openai-b", + Models: []string{"gpt-4o-mini"}, + UpstreamScheme: "https", + UpstreamHost: "b.example.com", + UpstreamPath: "/v1/chat", + AllowedGroupIDs: []string{"grp-eng"}, + }, + }}) + + in := newInputWithModelAndURL("gpt-4o-mini", "/v1/chat/completions") + in.UserGroups = []string{"grp-eng"} + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision) + + resolved, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID) + assert.Equal(t, "openai-b", resolved, + "longest-prefix path tiebreak still wins among group-authorised candidates") +} + +// TestRouter_AuthorisingGroups_EmitsIntersection pins that the router +// emits llm.authorising_groups containing only the intersection of the +// caller's UserGroups with the resolved route's AllowedGroupIDs — not +// every group the peer happens to be in. +func TestRouter_AuthorisingGroups_EmitsIntersection(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "openai-eng", + Models: []string{"gpt-4o-mini"}, + UpstreamScheme: "https", + UpstreamHost: "eng-openai.example.com", + AllowedGroupIDs: []string{"grp-eng", "grp-shared"}, + }}}) + + in := newInputWithModelAndGroups("gpt-4o-mini", + []string{"grp-eng", "grp-it", "grp-shared", "grp-oncall"}) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.Equal(t, middleware.DecisionAllow, out.Decision) + + csv, ok := metaValue(t, out.Metadata, middleware.KeyLLMAuthorisingGroups) + require.True(t, ok, "router must emit llm.authorising_groups on a match") + assert.Equal(t, "grp-eng,grp-shared", csv, + "only groups in BOTH UserGroups AND AllowedGroupIDs may appear; result must be sorted and unique") +} + +// TestRouter_EmptyModelsClaimsAnyModel pins that a route with no +// configured Models matches every model — used by gateway-style +// providers (LiteLLM, custom OpenAI-compatible endpoints) where the +// operator can't enumerate the upstream's model catalog in NetBird. +func TestRouter_EmptyModelsClaimsAnyModel(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "litellm", + Models: nil, // catch-all + UpstreamScheme: "https", + UpstreamHost: "litellm.example.com", + AllowedGroupIDs: []string{defaultTestGroup}, + }}}) + + out, err := mw.Invoke(context.Background(), newInputWithModel("gpt-5.5")) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, + "a route with empty Models must claim any model so gateway-style providers can route open-ended sets") + resolved, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID) + assert.Equal(t, "litellm", resolved) +} diff --git a/proxy/internal/middleware/builtin/llm_router/path_routed_test.go b/proxy/internal/middleware/builtin/llm_router/path_routed_test.go new file mode 100644 index 000000000..7dbd6b936 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_router/path_routed_test.go @@ -0,0 +1,159 @@ +package llm_router + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +// pathRoutedInput builds an Input mimicking the post-llm_request_parser state +// for a path-routed (Vertex/Bedrock) request: a request URL plus the model and +// (optionally) provider/vendor metadata the parser emits. +func pathRoutedInput(url, provider, model string) *middleware.Input { + md := []middleware.KV{{Key: middleware.KeyLLMModel, Value: model}} + if provider != "" { + md = append(md, middleware.KV{Key: middleware.KeyLLMProvider, Value: provider}) + } + return &middleware.Input{ + Slot: middleware.SlotOnRequest, + URL: url, + Metadata: md, + UserGroups: []string{defaultTestGroup}, + } +} + +func vertexRoute() ProviderRoute { + return ProviderRoute{ + ID: "vertex-prod", Vertex: true, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "europe-west1-aiplatform.googleapis.com", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer x", + } +} + +// A Vertex publisher with no parser surface (google/gemini emits no +// llm.provider) must be denied, not forwarded unmetered. +func TestRouter_VertexUnmeterablePublisherDenied(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{vertexRoute()}}) + in := pathRoutedInput( + "/v1/projects/p/locations/global/publishers/google/models/gemini-2.5-pro:generateContent", + "", // google -> request parser emits NO llm.provider + "gemini-2.5-pro", + ) + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionDeny, out.Decision, "unmeterable Vertex publisher must deny") + assert.Equal(t, 403, out.DenyStatus, "unmeterable deny is a 403") + require.NotNil(t, out.DenyReason) + assert.Equal(t, denyCodeUnmeterable, out.DenyReason.Code, "deny code must flag the unmeterable publisher") +} + +// A Vertex publisher with a parser surface (anthropic) is allowed. +func TestRouter_VertexMeterablePublisherAllowed(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{vertexRoute()}}) + in := pathRoutedInput( + "/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-4-5:rawPredict", + "anthropic", + "claude-sonnet-4-5", + ) + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "meterable Vertex publisher must allow") +} + +// A path-routed provider with an explicit Models list must reject models not in +// the list (the provider credential can't be used for unauthorised models). +func TestRouter_PathRoutedModelAllowlistEnforced(t *testing.T) { + route := ProviderRoute{ + ID: "bedrock-prod", Bedrock: true, + Models: []string{"anthropic.claude-sonnet-4-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer x", + } + mw := New(Config{Providers: []ProviderRoute{route}}) + + allowed := pathRoutedInput( + "/model/eu.anthropic.claude-sonnet-4-5-20250929-v1:0/invoke", + "bedrock", "anthropic.claude-sonnet-4-5", + ) + out, err := mw.Invoke(context.Background(), allowed) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "model in the allowlist must be served") + + denied := pathRoutedInput( + "/model/amazon.nova-pro-v1:0/invoke", + "bedrock", "amazon.nova-pro", + ) + out, err = mw.Invoke(context.Background(), denied) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionDeny, out.Decision, "model outside the allowlist must deny") + require.NotNil(t, out.DenyReason) + assert.Equal(t, denyCodeNotRoutable, out.DenyReason.Code, "unlisted model denies as not-routable") +} + +// A "/bedrock" gateway-namespace prefix routes the same as the native path and +// records the prefix on the rewrite so the proxy strips it before forwarding. +func TestRouter_BedrockNamespacePrefixStripped(t *testing.T) { + route := ProviderRoute{ + ID: "bedrock-prod", Bedrock: true, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer x", + } + mw := New(Config{Providers: []ProviderRoute{route}}) + + prefixed := pathRoutedInput( + "/bedrock/model/eu.anthropic.claude-sonnet-4-5-20250929-v1:0/invoke-with-response-stream", + "bedrock", "anthropic.claude-sonnet-4-5", + ) + out, err := mw.Invoke(context.Background(), prefixed) + require.NoError(t, err) + require.Equal(t, middleware.DecisionAllow, out.Decision, "prefixed Bedrock path must route") + require.NotNil(t, out.Mutations) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "/bedrock", out.Mutations.RewriteUpstream.StripPathPrefix, + "namespace prefix must be recorded so the proxy strips it before forwarding") + + native := pathRoutedInput( + "/model/eu.anthropic.claude-sonnet-4-5-20250929-v1:0/invoke", + "bedrock", "anthropic.claude-sonnet-4-5", + ) + out, err = mw.Invoke(context.Background(), native) + require.NoError(t, err) + require.Equal(t, middleware.DecisionAllow, out.Decision, "native Bedrock path must route") + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Empty(t, out.Mutations.RewriteUpstream.StripPathPrefix, + "native path carries no namespace prefix to strip") +} + +// A path-routed provider with no configured Models is catch-all: any model the +// credential can reach is served (preserves the zero-config behaviour). +func TestRouter_PathRoutedCatchAllServesAnyModel(t *testing.T) { + route := ProviderRoute{ + ID: "bedrock-catchall", Bedrock: true, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer x", + } + mw := New(Config{Providers: []ProviderRoute{route}}) + in := pathRoutedInput( + "/model/amazon.nova-pro-v1:0/invoke", + "bedrock", "amazon.nova-pro", + ) + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "catch-all path-routed provider serves any model") +}