Merge branch 'main' into reverse-proxy-allow-match-or

# Conflicts:
#	management/server/store/sql_store.go
#	management/server/store/sql_store_service_test.go
This commit is contained in:
Viktor Liu
2026-07-31 21:57:57 +02:00
216 changed files with 13000 additions and 4032 deletions

View File

@@ -3,20 +3,30 @@ package auth
import (
"context"
"net/netip"
"os"
"strings"
"sync"
"time"
log "github.com/sirupsen/logrus"
"golang.org/x/sync/singleflight"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/management/proto"
)
// tunnelCacheTTL caps how long a positive ValidateTunnelPeer result is
// reused before re-fetching from management. 5 minutes balances freshness
// against management load on busy mesh networks.
// tunnelCacheTTL is the default cap on how long a positive ValidateTunnelPeer
// result is reused before re-fetching from management. 5 minutes balances
// freshness against management load on busy mesh networks. Override it with
// envTunnelCacheTTL when an account needs authorization changes to take effect
// sooner (at the cost of more ValidateTunnelPeer RPCs).
const tunnelCacheTTL = 300 * time.Second
// envTunnelCacheTTL overrides tunnelCacheTTL. The value is a Go duration string
// (e.g. "30s", "2m"); an unset, unparseable, or non-positive value keeps the
// default.
const envTunnelCacheTTL = "NB_PROXY_TUNNEL_CACHE_TTL"
// tunnelCachePerAccount caps the number of cached identities per account.
// Bounded eviction avoids memory growth in pathological cases (huge peer
// roster, brief request bursts) while staying generous for normal use.
@@ -60,16 +70,35 @@ type accountBucket struct {
order []tunnelCacheKey
}
// newTunnelValidationCache constructs a cache with default TTL and bounds.
// newTunnelValidationCache constructs a cache with the configured TTL
// (envTunnelCacheTTL override or default) and default bounds.
func newTunnelValidationCache() *tunnelValidationCache {
return &tunnelValidationCache{
entries: make(map[types.AccountID]*accountBucket),
ttl: tunnelCacheTTL,
ttl: tunnelCacheTTLFromEnv(),
maxSize: tunnelCachePerAccount,
now: time.Now,
}
}
// tunnelCacheTTLFromEnv returns the tunnel-cache TTL, honoring the
// envTunnelCacheTTL override. The override must be a positive Go duration
// string (e.g. "30s", "2m"); anything unset, unparseable, or non-positive
// falls back to tunnelCacheTTL.
func tunnelCacheTTLFromEnv() time.Duration {
raw := strings.TrimSpace(os.Getenv(envTunnelCacheTTL))
if raw == "" {
return tunnelCacheTTL
}
d, err := time.ParseDuration(raw)
if err != nil || d <= 0 {
log.Warnf("ignoring invalid %s=%q (want a positive Go duration like 30s or 2m); using default %s",
envTunnelCacheTTL, raw, tunnelCacheTTL)
return tunnelCacheTTL
}
return d
}
// get returns a cached response for the key, or nil when missing or
// expired. Expired entries are evicted lazily on read.
func (c *tunnelValidationCache) get(key tunnelCacheKey) *proto.ValidateTunnelPeerResponse {

View File

@@ -169,3 +169,32 @@ func TestTunnelCache_BoundedSizeEvictsOldest(t *testing.T) {
assert.NotNil(t, cache.get(keys[1]), "second-newest must remain cached")
assert.NotNil(t, cache.get(keys[2]), "newest must remain cached")
}
func TestTunnelCacheTTLFromEnv(t *testing.T) {
t.Run("unset uses default", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "")
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
})
t.Run("valid duration overrides", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "45s")
assert.Equal(t, 45*time.Second, tunnelCacheTTLFromEnv())
})
t.Run("whitespace trimmed", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, " 2m ")
assert.Equal(t, 2*time.Minute, tunnelCacheTTLFromEnv())
})
t.Run("unparseable uses default", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "nonsense")
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
})
t.Run("non-positive uses default", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "0s")
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
t.Setenv(envTunnelCacheTTL, "-30s")
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
})
t.Run("constructor honors override", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "90s")
assert.Equal(t, 90*time.Second, newTunnelValidationCache().ttl)
})
}

View File

@@ -1,38 +0,0 @@
package llm
import (
"regexp"
"strings"
)
// 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+)?$`)
// 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 the inference-profile ARN's last segment likewise. It is the single
// source of truth shared by the request parser (which normalizes the request
// model from the URL path) and the router (which normalizes the operator's
// registered Bedrock model ids so both sides compare equal).
func NormalizeBedrockModel(modelID string) string {
m := modelID
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, "")
}

View File

@@ -1,23 +0,0 @@
package llm
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-haiku-4-5": "anthropic.claude-haiku-4-5",
"us.anthropic.claude-opus-4-8-20250101-v1:0": "anthropic.claude-opus-4-8",
"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",
// 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)
}
}

View File

@@ -1,59 +0,0 @@
# 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

View File

@@ -0,0 +1,21 @@
package llm
import (
sharedllm "github.com/netbirdio/netbird/shared/llm"
)
// 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. Thin delegate to the shared implementation
// (shared/llm), which management also uses at synthesis time so both sides of
// the pricing / routing contract normalize identically.
func NormalizeBedrockModel(modelID string) string {
return sharedllm.NormalizeBedrockModel(modelID)
}
// NormalizeVertexModel strips the "@version" suffix from a Vertex AI model id
// so it matches the catalog/pricing key. Thin delegate to shared/llm, kept
// beside NormalizeBedrockModel for the same contract reason.
func NormalizeVertexModel(modelID string) string {
return sharedllm.NormalizeVertexModel(modelID)
}

View File

@@ -1,65 +0,0 @@
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")
}

View File

@@ -1,288 +0,0 @@
# 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
# Kimi / Moonshot AI (kimi_api) — OpenAI-compatible /v1 endpoint. Moonshot
# reports cache hits OpenAI-style when present; cached input is 10% of
# input ($0.30 vs $3.00 per MTok). kimi-k3 is the only model the platform
# serves newer accounts (K2-era ids and kimi-latest 404), matching the
# management catalog.
kimi-k3:
input_per_1k: 0.003
output_per_1k: 0.015
cached_input_per_1k: 0.0003
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
# Kimi / Moonshot AI (kimi_api) via the Anthropic-compatible endpoint
# (/anthropic/v1/messages — the official Claude Code setup). Same rates
# as the OpenAI-shape entry above. "kimi-k3[1m]" is the model id some
# Claude Code guides set for the 1M-context alias; priced identically so
# cost metering doesn't silently skip those requests.
kimi-k3:
input_per_1k: 0.003
output_per_1k: 0.015
cache_read_per_1k: 0.0003
"kimi-k3[1m]":
input_per_1k: 0.003
output_per_1k: 0.015
cache_read_per_1k: 0.0003
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

View File

@@ -1,102 +1,30 @@
// 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 implements the pricing table and cost formula the
// cost_meter middleware uses to convert LLM token usage into a USD cost
// estimate. The table's content arrives from the management server inside
// cost_meter's middleware config (synthesized from the catalog plus the
// operator's stored per-provider prices) — the proxy carries no embedded
// price list. Price updates ride the ordinary mapping push: a chain
// rebuild constructs a fresh table, so there is nothing to reload.
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).
// configured", and cached tokens are billed at InputPer1K.
//
// 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.
// corresponding token bucket is billed at InputPer1K.
type Entry struct {
InputPer1K float64
OutputPer1K float64
@@ -105,33 +33,102 @@ type Entry struct {
CacheCreationPer1K float64
}
// Table is a provider-to-model pricing lookup. Instances are immutable once
// built and are swapped atomically by Loader.
// EntryJSON is the wire shape of a pricing entry inside cost_meter's
// middleware config. Field names are the management→proxy contract; the
// management synthesizer marshals the same names (its pricing.Entry).
type EntryJSON struct {
InputPer1K float64 `json:"input_per_1k"`
OutputPer1K float64 `json:"output_per_1k"`
CachedInputPer1K float64 `json:"cached_input_per_1k"`
CacheReadPer1K float64 `json:"cache_read_per_1k"`
CacheCreationPer1K float64 `json:"cache_creation_per_1k"`
}
// Table is a provider-surface-to-model pricing lookup. Instances are
// immutable once built; a mapping update builds a whole new middleware
// instance (and with it a new table) rather than mutating this one.
type Table struct {
entries map[string]map[string]Entry
}
// NewEntries validates and converts a wire-shape map (surface-or-record ->
// model -> rates) into the internal representation. Every rate must be a
// finite, non-negative USD amount; a violation is returned as an error so
// a corrupt config fails the chain build loudly instead of mispricing.
// Management validates the same constraints at its API boundary, so this
// is defense-in-depth. Nil input yields an empty (never-matching) map.
func NewEntries(raw map[string]map[string]EntryJSON) (map[string]map[string]Entry, error) {
out := make(map[string]map[string]Entry, len(raw))
for outer, models := range raw {
inner := make(map[string]Entry, len(models))
for model, e := range models {
for field, v := range map[string]float64{
"input_per_1k": e.InputPer1K,
"output_per_1k": e.OutputPer1K,
"cached_input_per_1k": e.CachedInputPer1K,
"cache_read_per_1k": e.CacheReadPer1K,
"cache_creation_per_1k": e.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", outer, model, field, v)
}
}
// EntryJSON and Entry are field-identical (tags aside), so a
// direct conversion carries all five rates.
inner[model] = Entry(e)
}
out[outer] = inner
}
return out, nil
}
// NewTable builds an immutable Table from the wire-shape defaults map.
// See NewEntries for validation semantics.
func NewTable(raw map[string]map[string]EntryJSON) (*Table, error) {
entries, err := NewEntries(raw)
if err != nil {
return nil, err
}
return &Table{entries: entries}, nil
}
// Lookup returns the entry for the given provider surface and model.
func (t *Table) Lookup(provider, model string) (Entry, bool) {
if t == nil {
return Entry{}, false
}
byModel, ok := t.entries[provider]
if !ok {
return Entry{}, false
}
e, ok := byModel[model]
return e, ok
}
// Has reports whether the provider/model pair is present in the table.
func (t *Table) Has(provider, model string) bool {
_, ok := t.Lookup(provider, model)
return ok
}
// 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) {
c, ok := t.Costs(provider, model, inTokens, outTokens, cachedInput, cacheCreation)
return c.TotalUSD, ok
}
// Costs returns the estimated USD cost split for the given token counts.
// The provider surface selects the cache formula; see EntryCosts.
func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (Costs, bool) {
entry, ok := t.Lookup(provider, model)
if !ok {
return Costs{}, false
}
return EntryCosts(entry, provider, inTokens, outTokens, cachedInput, cacheCreation), true
}
// Costs is a per-request cost split. The four per-bucket fields are the base
// of the breakdown — one per token bucket the provider bills separately — and
// the two aggregates are derived from them:
@@ -165,9 +162,25 @@ func newCosts(input, cachedInput, cacheCreation, output float64) Costs {
}
}
// Costs returns the estimated USD cost split for the given token counts, with
// the same semantics as Cost.
func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (Costs, bool) {
// EntryCosts computes the USD cost split for the given entry and token
// counts. The surface (the llm.provider value the request parser stamps)
// selects the cache formula; the entry may come from the surface-keyed
// defaults table or from a per-provider-record override — the math is
// identical either way.
//
// 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", "bedrock": 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 surfaces: cached and cacheCreation are ignored; cost is
// inTokens*InputPer1K + outTokens*OutputPer1K.
func EntryCosts(entry Entry, surface string, inTokens, outTokens, cachedInput, cacheCreation int64) Costs {
// Clamp negatives to zero before any pricing math so a malformed
// upstream count can never produce a negative cost.
if inTokens < 0 {
@@ -182,19 +195,8 @@ func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput,
if cacheCreation < 0 {
cacheCreation = 0
}
if t == nil {
return Costs{}, false
}
byModel, ok := t.entries[provider]
if !ok {
return Costs{}, false
}
entry, ok := byModel[model]
if !ok {
return Costs{}, false
}
output := (float64(outTokens) / 1000.0) * entry.OutputPer1K
switch provider {
switch surface {
case "openai":
// cachedInput is a subset of inTokens; clamp so a malformed
// upstream (cached > total) can't produce a negative remainder.
@@ -208,7 +210,7 @@ func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput,
}
nonCached := float64(inTokens-clamped) / 1000.0 * entry.InputPer1K
cached := float64(clamped) / 1000.0 * cachedRate
return newCosts(nonCached, cached, 0, output), true
return newCosts(nonCached, cached, 0, output)
case "anthropic", "bedrock":
// Bedrock-Anthropic returns the same additive cache buckets as
// first-party Anthropic; non-Anthropic Bedrock models simply report
@@ -224,266 +226,9 @@ func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput,
input := float64(inTokens) / 1000.0 * entry.InputPer1K
read := float64(cachedInput) / 1000.0 * readRate
create := float64(cacheCreation) / 1000.0 * createRate
return newCosts(input, read, create, output), true
return newCosts(input, read, create, output)
default:
input := float64(inTokens) / 1000.0 * entry.InputPer1K
return newCosts(input, 0, 0, output), true
return newCosts(input, 0, 0, output)
}
}
// 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
}

View File

@@ -1,20 +0,0 @@
//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)
}

View File

@@ -1,47 +1,13 @@
//go:build unix
package pricing
import (
"context"
"os"
"path/filepath"
"math"
"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
@@ -65,11 +31,9 @@ func TestCost_OpenAICachedSubsetDiscount(t *testing.T) {
"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).
// TestCost_OpenAICachedFallsBackToInputRate covers the fallback
// contract: when CachedInputPer1K is unset (zero), cached tokens bill
// at the regular input rate.
func TestCost_OpenAICachedFallsBackToInputRate(t *testing.T) {
tbl := &Table{entries: map[string]map[string]Entry{
"openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01}},
@@ -78,7 +42,7 @@ func TestCost_OpenAICachedFallsBackToInputRate(t *testing.T) {
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")
"absent cached_input_per_1k rate must fall back to input_per_1k")
}
// TestCost_OpenAIClampsCachedToInputCount is the defensive guard
@@ -100,10 +64,7 @@ func TestCost_OpenAIClampsCachedToInputCount(t *testing.T) {
// 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).
// configured rate.
func TestCost_AnthropicCacheReadAndCreationAreAdditive(t *testing.T) {
tbl := &Table{entries: map[string]map[string]Entry{
"anthropic": {"claude-sonnet": {
@@ -125,11 +86,9 @@ func TestCost_AnthropicCacheReadAndCreationAreAdditive(t *testing.T) {
"each Anthropic input bucket must bill at its own configured rate")
}
// TestCost_AnthropicCacheRatesFallBackToInput covers the no-opt-in
// TestCost_AnthropicCacheRatesFallBackToInput covers the no-rate
// 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.
// cache tokens bill at the regular input rate.
func TestCost_AnthropicCacheRatesFallBackToInput(t *testing.T) {
tbl := &Table{entries: map[string]map[string]Entry{
"anthropic": {"claude-sonnet": {InputPer1K: 0.003, OutputPer1K: 0.015}},
@@ -139,259 +98,39 @@ func TestCost_AnthropicCacheRatesFallBackToInput(t *testing.T) {
// 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")
"absent cache rates must fall back to input_per_1k")
}
func TestNewLoader_UnknownModel(t *testing.T) {
base := t.TempDir()
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
// TestEntryCosts_SurfaceSelectsFormula pins that the formula branches on
// the SURFACE, not on which table the entry came from: the same entry
// bills a subset carve-out on "openai", additive buckets on
// "anthropic"/"bedrock", and ignores cache counts everywhere else. This
// is what keeps per-provider-record entries (looked up by record id)
// mathematically identical to defaults-table entries.
func TestEntryCosts_SurfaceSelectsFormula(t *testing.T) {
e := Entry{InputPer1K: 0.002, OutputPer1K: 0.01, CachedInputPer1K: 0.001, CacheReadPer1K: 0.0002, CacheCreationPer1K: 0.0025}
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
require.NoError(t, err)
openai := EntryCosts(e, "openai", 1000, 0, 400, 300)
assert.InDelta(t, (600.0/1000.0)*0.002+(400.0/1000.0)*0.001, openai.TotalUSD, 1e-12,
"openai: cached is a subset, cacheCreation ignored")
_, ok := l.Get().Cost("openai", "fantasy-model", 10, 10, 0, 0)
assert.False(t, ok, "unknown model returns ok=false")
anthropic := EntryCosts(e, "anthropic", 1000, 0, 400, 300)
assert.InDelta(t, 0.002+(400.0/1000.0)*0.0002+(300.0/1000.0)*0.0025, anthropic.TotalUSD, 1e-12,
"anthropic: cache buckets are additive")
_, ok = l.Get().Cost("cohere", "anything", 10, 10, 0, 0)
assert.False(t, ok, "unknown provider returns ok=false")
bedrock := EntryCosts(e, "bedrock", 1000, 0, 400, 300)
assert.InDelta(t, anthropic.TotalUSD, bedrock.TotalUSD, 1e-12, "bedrock shares the anthropic formula")
other := EntryCosts(e, "gemini", 1000, 0, 400, 300)
assert.InDelta(t, 0.002, other.TotalUSD, 1e-12, "unknown surface: cache counts ignored")
}
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")
// TestEntryCosts_ClampsNegativeTokens: malformed upstream counts must
// never produce a negative cost.
func TestEntryCosts_ClampsNegativeTokens(t *testing.T) {
e := Entry{InputPer1K: 0.002, OutputPer1K: 0.01}
c := EntryCosts(e, "openai", -50, -10, -5, -3)
assert.Zero(t, c.TotalUSD, "all-negative counts clamp to zero cost")
}
func TestTableCost_NilSafe(t *testing.T) {
@@ -402,31 +141,37 @@ func TestTableCost_NilSafe(t *testing.T) {
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] = ' '
func TestNewTable_ValidatesRates(t *testing.T) {
good := map[string]map[string]EntryJSON{
"openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01, CachedInputPer1K: 0.00125}},
}
bigComment[0] = '#'
bigComment[len(bigComment)-1] = '\n'
payload := append([]byte(header), bigComment...)
require.NoError(t, os.WriteFile(target, payload, 0o600))
tbl, err := NewTable(good)
require.NoError(t, err)
cost, ok := tbl.Cost("openai", "gpt-4o", 1000, 1000, 0, 0)
require.True(t, ok, "entry survives the wire conversion")
assert.InDelta(t, 0.0125, cost, 1e-9)
_, 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")
_, ok = tbl.Cost("openai", "unknown-model", 1, 1, 0, 0)
assert.False(t, ok, "unknown model misses")
for name, bad := range map[string]EntryJSON{
"negative input": {InputPer1K: -1, OutputPer1K: 0.01},
"NaN output": {InputPer1K: 0.01, OutputPer1K: math.NaN()},
"Inf cache read": {InputPer1K: 0.01, OutputPer1K: 0.01, CacheReadPer1K: math.Inf(1)},
"negative cached": {InputPer1K: 0.01, OutputPer1K: 0.01, CachedInputPer1K: -0.001},
} {
_, err := NewTable(map[string]map[string]EntryJSON{"openai": {"m": bad}})
assert.Error(t, err, "case %q must be rejected so a corrupt config fails the chain build instead of mispricing", name)
}
}
func TestNewTable_NilAndEmpty(t *testing.T) {
tbl, err := NewTable(nil)
require.NoError(t, err, "nil map builds an empty (never-matching) table")
_, ok := tbl.Cost("openai", "gpt-4o", 1, 1, 0, 0)
assert.False(t, ok, "empty table prices nothing")
entries, err := NewEntries(nil)
require.NoError(t, err)
assert.Empty(t, entries, "nil in, empty (never-matching) map out for the per-record map")
}

View File

@@ -1,68 +0,0 @@
//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
}

View File

@@ -36,15 +36,13 @@ 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.
// 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
@@ -58,12 +56,11 @@ var (
// 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) {
func Configure(ctx context.Context, 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,

View File

@@ -12,6 +12,7 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
mgmtpricing "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/pricing"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/cost_meter"
@@ -19,17 +20,20 @@ import (
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_response_parser"
)
// Drives the real pipeline (llm_request_parser → llm_response_parser → cost_meter) on the embedded default pricing
// table and asserts exact USD amounts hardcoded from the vendors' published prices, including the cache split.
// Drives the real pipeline (llm_request_parser → llm_response_parser → cost_meter) on the REAL default pricing
// table management ships (mgmtpricing.DefaultTable, catalog-derived) and asserts exact USD amounts hardcoded from
// the vendors' published prices, including the cache split. This is the cross-stack pricing contract test: the
// management-side Entry JSON must decode into the proxy-side table and produce these exact costs.
func TestCostCalculation_ProviderMatrix(t *testing.T) {
// Empty data dir → embedded defaults, like a proxy with no pricing override.
builtin.Configure(context.Background(), t.TempDir(), nil, nil, nil)
builtin.Configure(context.Background(), nil, nil, nil)
reqMW, err := llm_request_parser.Factory{}.New(nil)
require.NoError(t, err, "build llm_request_parser")
respMW, err := llm_response_parser.Factory{}.New(nil)
require.NoError(t, err, "build llm_response_parser")
costMW, err := cost_meter.Factory{}.New(nil)
costCfgJSON, err := json.Marshal(map[string]any{"pricing": map[string]any{"defaults": mgmtpricing.DefaultTable()}})
require.NoError(t, err, "marshal management default table into cost_meter config")
costMW, err := cost_meter.Factory{}.New(costCfgJSON)
require.NoError(t, err, "build cost_meter")
t.Cleanup(func() { _ = costMW.Close() })

View File

@@ -2,7 +2,6 @@ package cost_meter
import (
"bytes"
"context"
"encoding/json"
"fmt"
@@ -11,16 +10,27 @@ import (
"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.
// Config is the on-wire configuration for the middleware, synthesized by
// management (buildCostMeterConfigJSON). The proxy has no embedded price
// list: this payload is the only pricing source, and updates arrive as
// ordinary mapping pushes that rebuild the chain (and with it this
// middleware instance) — no per-request fetches, no reload loops.
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"`
Pricing *PricingConfig `json:"pricing"`
}
// PricingConfig carries the full pricing table:
// - Defaults: parser surface ("openai"/"anthropic"/"bedrock") ->
// normalized model id -> rates, matched against llm.provider +
// llm.model.
// - Providers: provider record id -> normalized model id -> rates,
// matched against the llm.resolved_provider_id metadata llm_router
// stamps. Entries arrive fully materialized (management folds default
// cache rates in at synth time), so lookup order is simply
// per-record first, defaults second.
type PricingConfig struct {
Defaults map[string]map[string]pricing.EntryJSON `json:"defaults"`
Providers map[string]map[string]pricing.EntryJSON `json:"providers"`
}
// Factory builds cost_meter instances from raw config bytes.
@@ -29,45 +39,45 @@ 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.
// New constructs a middleware instance. Empty, null, and {} configs are
// accepted for backward compatibility with a management server that
// predates config-delivered pricing — the instance then skips every cost
// computation (unknown_model) and a warning is logged once at build time.
// Non-empty rawConfig that fails to unmarshal, or a table carrying a
// non-finite / negative rate, is rejected so misconfigurations surface at
// chain build time.
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
if cfg.Pricing == nil {
if logger := builtin.Context().Logger; logger != nil {
logger.Warnf("cost_meter: no pricing table in middleware config; management predates config-delivered pricing — every request will record cost.skipped=unknown_model ($0)")
}
return newMiddleware(mustEmptyTable(), nil), nil
}
loader, err := pricing.NewLoader(fctx.DataDir, pricingPath, ID, nil)
defaults, err := pricing.NewTable(cfg.Pricing.Defaults)
if err != nil {
return nil, fmt.Errorf("init pricing loader: %w", err)
return nil, fmt.Errorf("cost_meter pricing defaults: %w", err)
}
cancel := startReloader(fctx.Context, loader)
return newMiddleware(loader, cancel), nil
perRecord, err := pricing.NewEntries(cfg.Pricing.Providers)
if err != nil {
return nil, fmt.Errorf("cost_meter per-provider pricing: %w", err)
}
return newMiddleware(defaults, perRecord), 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
// mustEmptyTable returns a valid empty table. NewTable on a nil map cannot
// fail; the panic guard documents that invariant.
func mustEmptyTable() *pricing.Table {
t, err := pricing.NewTable(nil)
if err != nil {
panic(fmt.Sprintf("cost_meter: empty pricing table must build: %v", err))
}
cctx, cancel := context.WithCancel(ctx)
go loader.Reload(cctx)
return cancel
return t
}
// decodeConfig accepts empty, null, and {} configs, returning a

View File

@@ -1,7 +1,9 @@
// 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.
// per-request USD cost estimate. Pricing arrives from management inside
// the middleware config: a per-provider-record table (the operator's
// stored prices, matched via llm.resolved_provider_id) consulted first,
// then the surface-keyed defaults table.
package cost_meter
import (
@@ -17,7 +19,9 @@ import (
const ID = "cost_meter"
// Version is the implementation version emitted via the spec merge.
const Version = "1.0.0"
// 1.1.0: pricing is config-delivered (defaults + per-provider-record
// entries) instead of proxy-embedded.
const Version = "1.1.0"
// Skip reasons emitted under KeyCostSkipped. The set is closed; the
// dashboard surfaces these verbatim.
@@ -42,19 +46,21 @@ var metadataKeys = []string{
}
// Middleware computes a per-response cost estimate from the token
// counts emitted upstream by llm_response_parser.
// counts emitted upstream by llm_response_parser. Both tables are
// immutable — a pricing change arrives as a mapping push that rebuilds
// the chain with a fresh instance.
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
// defaults is the surface-keyed table (llm.provider x llm.model).
defaults *pricing.Table
// perRecord is keyed by provider record id (llm.resolved_provider_id)
// then normalized model id; entries arrive fully materialized from
// management. Consulted before defaults. May be nil.
perRecord map[string]map[string]pricing.Entry
}
// 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}
// newMiddleware constructs a Middleware over the given pricing tables.
func newMiddleware(defaults *pricing.Table, perRecord map[string]map[string]pricing.Entry) *Middleware {
return &Middleware{defaults: defaults, perRecord: perRecord}
}
// ID returns the registry identifier.
@@ -79,16 +85,9 @@ func (m *Middleware) MetadataKeys() []string {
// 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
}
// Close releases resources owned by the middleware. Stateless — the
// pricing tables are plain maps owned by this instance.
func (m *Middleware) Close() error { return nil }
// Invoke reads provider, model, and token metadata, looks up pricing,
// and emits either KeyCostUSDTotal or KeyCostSkipped. The decision is
@@ -144,8 +143,7 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
return out, nil
}
table := m.loader.Get()
costs, ok := table.Costs(provider, model, inTokens, outTokens, cachedTokens, cacheCreationTokens)
costs, ok := m.lookupCosts(in.Metadata, provider, model, inTokens, outTokens, cachedTokens, cacheCreationTokens)
if !ok {
out.Metadata = skip(skipUnknownModel)
return out, nil
@@ -164,6 +162,26 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
return out, nil
}
// lookupCosts resolves the price for this request and computes the cost
// split. Resolution order:
//
// 1. Per-provider-record entry: the operator's stored price for the
// provider route that served the request, keyed by the
// llm.resolved_provider_id metadata llm_router stamped on the allow
// path. Absent metadata (e.g. no router in the chain) skips this tier.
// 2. Surface defaults: the catalog-derived table keyed by llm.provider.
//
// The surface always selects the cache formula — a per-record entry for an
// Anthropic route still bills its cache buckets additively.
func (m *Middleware) lookupCosts(md []middleware.KV, surface, model string, inTokens, outTokens, cachedTokens, cacheCreationTokens int64) (pricing.Costs, bool) {
if recordID := lookupKV(md, middleware.KeyLLMResolvedProviderID); recordID != "" {
if entry, ok := m.perRecord[recordID][model]; ok {
return pricing.EntryCosts(entry, surface, inTokens, outTokens, cachedTokens, cacheCreationTokens), true
}
}
return m.defaults.Costs(surface, model, inTokens, outTokens, cachedTokens, cacheCreationTokens)
}
// usd renders a cost as the fixed-precision string every cost.usd_* key
// carries, so the per-bucket values and the aggregates round identically.
//

View File

@@ -3,39 +3,50 @@ 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/llm/pricing"
"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 {
// fixtureConfig mirrors what management's buildCostMeterConfigJSON ships:
// a surface-keyed defaults table. Rates match the retired YAML fixture so
// every cost assertion below is byte-identical to the pre-feature values.
func fixtureConfig(t *testing.T) []byte {
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
raw, err := json.Marshal(Config{Pricing: &PricingConfig{
Defaults: map[string]map[string]pricing.EntryJSON{
"openai": {
"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01},
"gpt-4o-mini": {InputPer1K: 0.00015, OutputPer1K: 0.0006},
},
"anthropic": {
"claude-sonnet-4-5": {InputPer1K: 0.003, OutputPer1K: 0.015},
},
},
}})
require.NoError(t, err)
return raw
}
// fixtureConfigWithCache adds the cache-rate fields.
func fixtureConfigWithCache(t *testing.T) []byte {
t.Helper()
raw, err := json.Marshal(Config{Pricing: &PricingConfig{
Defaults: map[string]map[string]pricing.EntryJSON{
"openai": {
"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01, CachedInputPer1K: 0.00125},
},
"anthropic": {
"claude-sonnet-4-5": {InputPer1K: 0.003, OutputPer1K: 0.015, CacheReadPer1K: 0.0003, CacheCreationPer1K: 0.00375},
},
},
}})
require.NoError(t, err)
return raw
}
func metaValue(t *testing.T, kvs []middleware.KV, key string) (string, bool) {
@@ -56,8 +67,7 @@ func buildMiddleware(t *testing.T, raw []byte) middleware.Middleware {
}
func TestMiddleware_StaticSurface(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
mw := buildMiddleware(t, fixtureConfig(t))
assert.Equal(t, ID, mw.ID(), "ID must match the registered constant")
assert.Equal(t, Version, mw.Version(), "Version must match the constant")
@@ -79,8 +89,10 @@ func TestMiddleware_StaticSurface(t *testing.T) {
assert.Equal(t, expected, keys, "metadata key allowlist must match the spec")
}
// TestFactory_AcceptsEmptyAndJSONConfig: empty/null/{} configs are what an
// old management (pre config-delivered pricing) sends — they must build a
// working (all-skip) instance, never fail the chain.
func TestFactory_AcceptsEmptyAndJSONConfig(t *testing.T) {
configureBuiltin(t)
cases := [][]byte{nil, {}, []byte("null"), []byte("{}"), []byte(" ")}
for _, raw := range cases {
mw, err := Factory{}.New(raw)
@@ -90,15 +102,57 @@ func TestFactory_AcceptsEmptyAndJSONConfig(t *testing.T) {
}
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)
// TestFactory_RejectsInvalidRates: a non-finite or negative rate anywhere
// in the table fails the chain build (defense-in-depth behind management's
// API validation) rather than silently mispricing.
func TestFactory_RejectsInvalidRates(t *testing.T) {
raw, err := json.Marshal(Config{Pricing: &PricingConfig{
Defaults: map[string]map[string]pricing.EntryJSON{
"openai": {"gpt-4o": {InputPer1K: -0.0025, OutputPer1K: 0.01}},
},
}})
require.NoError(t, err)
mw, err := Factory{}.New(raw)
require.Error(t, err, "negative rate must fail the build")
assert.Nil(t, mw)
raw, err = json.Marshal(Config{Pricing: &PricingConfig{
Providers: map[string]map[string]pricing.EntryJSON{
"prov-1": {"m": {InputPer1K: 0.01, OutputPer1K: 0.01, CacheReadPer1K: -1}},
},
}})
require.NoError(t, err)
_, err = Factory{}.New(raw)
require.Error(t, err, "per-record tables validate too")
}
// TestFactory_NilPricingSkipsEverything is the version-skew contract: a
// new proxy under an old management ({} config) must build, allow, and
// skip with unknown_model — degraded but never broken.
func TestFactory_NilPricingSkipsEverything(t *testing.T) {
mw := buildMiddleware(t, []byte("{}"))
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: "1000"},
},
})
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision, "cost_meter always allows")
value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
require.True(t, ok, "no pricing table means every request skips")
assert.Equal(t, skipUnknownModel, value)
}
func TestFactory_ConfigDefaultsPriceRequests(t *testing.T) {
mw := buildMiddleware(t, fixtureConfig(t))
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
@@ -116,34 +170,78 @@ func TestFactory_DefaultPricingPathLoadsFixture(t *testing.T) {
assert.Equal(t, "0.000750000", value, "0.00015 + 0.0006 per 1k tokens, 9-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"})
// TestInvoke_PerRecordEntryBeatsDefaults: when llm_router resolved a
// provider record whose operator pinned a price for the model, that price
// wins over the surface default.
func TestInvoke_PerRecordEntryBeatsDefaults(t *testing.T) {
raw, err := json.Marshal(Config{Pricing: &PricingConfig{
Defaults: map[string]map[string]pricing.EntryJSON{
"openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01}},
},
Providers: map[string]map[string]pricing.EntryJSON{
"prov-azure": {"gpt-4o": {InputPer1K: 0.005, OutputPer1K: 0.02}},
},
}})
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.KeyLLMResolvedProviderID, Value: "prov-azure"},
{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 with custom pricing path")
assert.Equal(t, "0.015000000", value, "2*0.0025 + 1*0.01 = 0.015 with 9-decimal format")
require.True(t, ok)
assert.Equal(t, "0.025000000", value, "operator's per-record price (0.005+0.02) wins over the default (0.0025+0.01)")
}
func TestInvoke_ComputesCostForKnownModel(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
// TestInvoke_PerRecordMissFallsBackToDefaults: a resolved record with no
// entry for this model (or no entries at all) falls through to the
// surface defaults — gateway providers rely on exactly this.
func TestInvoke_PerRecordMissFallsBackToDefaults(t *testing.T) {
raw, err := json.Marshal(Config{Pricing: &PricingConfig{
Defaults: map[string]map[string]pricing.EntryJSON{
"openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01}},
},
Providers: map[string]map[string]pricing.EntryJSON{
"prov-1": {"some-other-model": {InputPer1K: 1, OutputPer1K: 1}},
},
}})
require.NoError(t, err)
mw := buildMiddleware(t, raw)
for name, recordID := range map[string]string{
"record with other models": "prov-1",
"record with no entries": "prov-gateway",
} {
t.Run(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.KeyLLMResolvedProviderID, Value: recordID},
{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, "per-record miss must fall back to the surface default, not skip")
assert.Equal(t, "0.012500000", value, "default rates apply")
})
}
}
// TestInvoke_NoResolvedProviderIDUsesDefaults: metadata without a
// resolved provider id (router denied, or a chain without llm_router)
// prices from the defaults table directly.
func TestInvoke_NoResolvedProviderIDUsesDefaults(t *testing.T) {
mw := buildMiddleware(t, fixtureConfig(t))
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
{Key: middleware.KeyLLMProvider, Value: "anthropic"},
@@ -153,17 +251,15 @@ func TestInvoke_ComputesCostForKnownModel(t *testing.T) {
},
})
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "cost.usd_total must be emitted")
require.True(t, ok)
assert.Equal(t, "0.018000000", value, "0.003 + 0.015 = 0.018 with 9-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)
mw := buildMiddleware(t, fixtureConfig(t))
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
@@ -179,8 +275,7 @@ func TestInvoke_MissingProvider(t *testing.T) {
}
func TestInvoke_MissingModel(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
mw := buildMiddleware(t, fixtureConfig(t))
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
@@ -196,8 +291,7 @@ func TestInvoke_MissingModel(t *testing.T) {
}
func TestInvoke_MissingTokens(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
mw := buildMiddleware(t, fixtureConfig(t))
cases := []struct {
name string
@@ -240,8 +334,7 @@ func TestInvoke_MissingTokens(t *testing.T) {
}
func TestInvoke_UnparseableTokens(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
mw := buildMiddleware(t, fixtureConfig(t))
cases := []struct {
name string
@@ -271,8 +364,7 @@ func TestInvoke_UnparseableTokens(t *testing.T) {
}
func TestInvoke_ZeroTokens(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
mw := buildMiddleware(t, fixtureConfig(t))
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
@@ -291,8 +383,7 @@ func TestInvoke_ZeroTokens(t *testing.T) {
}
func TestInvoke_UnknownModel(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
mw := buildMiddleware(t, fixtureConfig(t))
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
@@ -309,8 +400,7 @@ func TestInvoke_UnknownModel(t *testing.T) {
}
func TestInvoke_NilInput(t *testing.T) {
configureBuiltin(t)
mw := buildMiddleware(t, nil)
mw := buildMiddleware(t, fixtureConfig(t))
out, err := mw.Invoke(context.Background(), nil)
require.NoError(t, err)
@@ -319,36 +409,12 @@ func TestInvoke_NilInput(t *testing.T) {
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)
mw := buildMiddleware(t, fixtureConfigWithCache(t))
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
@@ -390,8 +456,7 @@ func TestInvoke_OpenAICachedSubsetDiscount(t *testing.T) {
// 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)
mw := buildMiddleware(t, fixtureConfigWithCache(t))
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
@@ -429,8 +494,37 @@ func TestInvoke_AnthropicCacheBucketsAdditive(t *testing.T) {
"output bucket bills 200 tokens at 0.015/1k")
}
// TestInvoke_PerRecordEntryUsesSurfaceFormula: a per-record entry for an
// anthropic-surface request must bill its cache buckets additively — the
// formula follows llm.provider, not which table the entry came from.
func TestInvoke_PerRecordEntryUsesSurfaceFormula(t *testing.T) {
raw, err := json.Marshal(Config{Pricing: &PricingConfig{
Providers: map[string]map[string]pricing.EntryJSON{
"prov-ant": {"claude-sonnet-4-5": {InputPer1K: 0.003, OutputPer1K: 0.015, CacheReadPer1K: 0.0003, CacheCreationPer1K: 0.00375}},
},
}})
require.NoError(t, err)
mw := buildMiddleware(t, raw)
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.KeyLLMResolvedProviderID, Value: "prov-ant"},
{Key: middleware.KeyLLMInputTokens, Value: "256"},
{Key: middleware.KeyLLMOutputTokens, Value: "200"},
{Key: middleware.KeyLLMCachedInputTokens, Value: "768"},
{Key: middleware.KeyLLMCacheCreationTokens, Value: "512"},
},
})
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok)
assert.Equal(t, "0.005918400", value, "identical math to the defaults-table entry with the same rates")
}
// assertBucket asserts one per-bucket cost key carries the expected
// 6-decimal value.
// 9-decimal value.
func assertBucket(t *testing.T, md []middleware.KV, key, want, msg string) {
t.Helper()
got, ok := metaValue(t, md, key)
@@ -439,13 +533,10 @@ func assertBucket(t *testing.T, md []middleware.KV, key, want, msg string) {
}
// 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.
// no-cache-metadata path: with no cached keys emitted, the meter must
// produce exactly the input+output cost.
func TestInvoke_CachedTokensAbsentFallsBackToBaseFormula(t *testing.T) {
configureBuiltinWithCacheRates(t)
mw := buildMiddleware(t, nil)
mw := buildMiddleware(t, fixtureConfigWithCache(t))
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
@@ -460,7 +551,7 @@ func TestInvoke_CachedTokensAbsentFallsBackToBaseFormula(t *testing.T) {
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.007500000", value, "no cached metadata = same cost as before the feature landed")
assert.Equal(t, "0.007500000", value, "no cached metadata = plain input+output cost")
}
// TestInvoke_UnparseableCachedTokensSkippedSilently proves the
@@ -469,8 +560,7 @@ func TestInvoke_CachedTokensAbsentFallsBackToBaseFormula(t *testing.T) {
// 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)
mw := buildMiddleware(t, fixtureConfigWithCache(t))
out, err := mw.Invoke(context.Background(), &middleware.Input{
Metadata: []middleware.KV{
@@ -487,22 +577,10 @@ func TestInvoke_UnparseableCachedTokensSkippedSilently(t *testing.T) {
assert.Equal(t, "0.007500000", 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.
// TestMiddleware_CloseNilSafe confirms Close is a no-op (no panic) even
// 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")
require.NoError(t, newMiddleware(nil, nil).Close(), "Close must be a no-op")
var m *Middleware
require.NoError(t, m.Close(), "nil-receiver Close must be safe")
}

View File

@@ -10,11 +10,15 @@ import (
)
// Config is the JSON-decoded shape accepted by the factory. The
// runtime path consumes the normalised allowlist; raw config is not
// runtime path consumes the normalised allowlists; raw config is not
// retained beyond construction.
type Config struct {
ModelAllowlist []string `json:"model_allowlist"`
PromptCapture PromptCapture `json:"prompt_capture"`
// ProviderAllowlists maps a resolved provider id (KeyLLMResolvedProviderID) to
// its model allowlist. A provider present is restricted to those models; one
// absent is unrestricted. Kept per-provider so one provider's list can't leak
// onto another.
ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"`
PromptCapture PromptCapture `json:"prompt_capture"`
}
// PromptCapture toggles the optional prompt capture + redaction step
@@ -54,21 +58,28 @@ func isEmptyJSON(raw []byte) bool {
return false
}
// normaliseConfig lowercases and trims allowlist entries so the runtime
// match is case-insensitive. Empty entries are dropped.
// normaliseConfig lowercases and trims allowlist entries for case-insensitive
// matching; empty entries drop. A provider whose entries all drop keeps an empty
// (non-nil) list — "deny every model" — distinct from an absent provider
// (unrestricted).
func normaliseConfig(cfg Config) Config {
if len(cfg.ModelAllowlist) == 0 {
if len(cfg.ProviderAllowlists) == 0 {
cfg.ProviderAllowlists = nil
return cfg
}
cleaned := make([]string, 0, len(cfg.ModelAllowlist))
for _, entry := range cfg.ModelAllowlist {
n := normaliseModel(entry)
if n == "" {
continue
cleaned := make(map[string][]string, len(cfg.ProviderAllowlists))
for provider, models := range cfg.ProviderAllowlists {
list := make([]string, 0, len(models))
for _, entry := range models {
n := normaliseModel(entry)
if n == "" {
continue
}
list = append(list, n)
}
cleaned = append(cleaned, n)
cleaned[provider] = list
}
cfg.ModelAllowlist = cleaned
cfg.ProviderAllowlists = cleaned
return cfg
}

View File

@@ -83,8 +83,9 @@ func (m *Middleware) MutationsSupported() bool { return false }
// 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)
providerID, _ := lookupMetadata(in.Metadata, middleware.KeyLLMResolvedProviderID)
if denial := m.evaluateAllowlist(model, modelPresent); denial != nil {
if denial := m.evaluateAllowlist(providerID, model, modelPresent); denial != nil {
return denial, nil
}
@@ -110,20 +111,32 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
// 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 {
// evaluateAllowlist denies when the resolved provider's allowlist rejects the
// model; nil means proceed. Scoped to the provider llm_router resolved, so an
// unrestricted provider (absent from config) is never caught by another's list.
func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bool) *middleware.Output {
if len(m.cfg.ProviderAllowlists) == 0 {
return nil
}
// Fail closed: with an allowlist configured, a request whose model the
// upstream parser could not extract (absent or empty) must be denied rather
// than allowed. This is what enforces the allowlist for URL/path-routed
// providers (Bedrock, Vertex, ...) whose model lives outside the JSON body.
// Restrictions exist but the resolved provider is unknown, so we can't tell
// if this request targets a restricted provider — fail closed. llm_router
// normally stamps the provider first, so this is a defensive guard.
if providerID == "" {
return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
}
allowlist, restricted := m.cfg.ProviderAllowlists[providerID]
if !restricted {
// This provider has no allowlist (some authorising policy left it
// unrestricted); management owns any per-policy/group decision.
return nil
}
// Fail closed: with an allowlist in effect for this provider, a request whose
// model the parser couldn't extract (absent/empty) is denied. This enforces
// the allowlist for path-routed providers (Bedrock, Vertex) with no body model.
if !modelPresent || normaliseModel(model) == "" {
return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
}
if m.modelInAllowlist(model) {
if modelInAllowlist(allowlist, model) {
return nil
}
return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel)
@@ -151,14 +164,15 @@ func denyModel(model, code, message, reason string) *middleware.Output {
}
}
// modelInAllowlist reports whether the model matches any allowlist
// entry under the case-insensitive, trim-tolerant comparison rule.
func (m *Middleware) modelInAllowlist(model string) bool {
// modelInAllowlist reports whether the model matches any entry in the supplied
// (already-normalised) allowlist under the case-insensitive, trim-tolerant
// comparison rule.
func modelInAllowlist(allowlist []string, model string) bool {
normalised := normaliseModel(model)
if normalised == "" {
return false
}
for _, allowed := range m.cfg.ModelAllowlist {
for _, allowed := range allowlist {
if allowed == normalised {
return true
}

View File

@@ -26,6 +26,25 @@ func newInput(meta ...middleware.KV) *middleware.Input {
return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: meta}
}
const (
testProvider = "prov-1"
otherProvider = "prov-2"
)
// providerCfg builds a Config restricting testProvider to the given models.
func providerCfg(models ...string) Config {
return Config{ProviderAllowlists: map[string][]string{testProvider: models}}
}
// newInputProvider builds an input that carries a resolved provider id (as
// llm_router would stamp) plus any extra metadata.
func newInputProvider(provider string, meta ...middleware.KV) *middleware.Input {
all := make([]middleware.KV, 0, len(meta)+1)
all = append(all, middleware.KV{Key: middleware.KeyLLMResolvedProviderID, Value: provider})
all = append(all, meta...)
return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: all}
}
func TestMiddlewareIdentity(t *testing.T) {
mw := New(Config{})
assert.Equal(t, ID, mw.ID(), "middleware ID must be llm_guardrail")
@@ -47,12 +66,12 @@ func TestMiddlewareIdentity(t *testing.T) {
func TestAllowlistEmptyAllowsAnyModel(t *testing.T) {
mw := New(Config{})
out, err := mw.Invoke(context.Background(), newInput(
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
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")
assert.Equal(t, middleware.DecisionAllow, out.Decision, "no provider allowlists 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")
@@ -62,8 +81,8 @@ func TestAllowlistEmptyAllowsAnyModel(t *testing.T) {
}
func TestAllowlistMatchAllows(t *testing.T) {
mw := New(Config{ModelAllowlist: []string{"gpt-4o", "claude-opus-4"}})
out, err := mw.Invoke(context.Background(), newInput(
mw := New(providerCfg("gpt-4o", "claude-opus-4"))
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
))
require.NoError(t, err)
@@ -71,8 +90,8 @@ func TestAllowlistMatchAllows(t *testing.T) {
}
func TestAllowlistMissDenies(t *testing.T) {
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
out, err := mw.Invoke(context.Background(), newInput(
mw := New(providerCfg("gpt-4o"))
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
))
require.NoError(t, err)
@@ -91,10 +110,10 @@ func TestAllowlistMissDenies(t *testing.T) {
}
func TestAllowlistCaseInsensitive(t *testing.T) {
mw := New(Config{ModelAllowlist: []string{" GPT-4o ", "Claude-OPUS-4"}})
mw := New(providerCfg(" 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(
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: model},
))
require.NoError(t, err)
@@ -103,14 +122,15 @@ func TestAllowlistCaseInsensitive(t *testing.T) {
}
func TestAllowlistMissingModelKeyDenies(t *testing.T) {
// Fail closed: with an allowlist configured, a request whose model the
// parser could not extract (URL/path-routed providers such as Bedrock or
// Vertex whose shape wasn't recognised) must be denied, not allowed.
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
out, err := mw.Invoke(context.Background(), newInput())
// Fail closed: with an allowlist in effect for the resolved provider, a
// request whose model the parser could not extract (URL/path-routed
// providers such as Bedrock or Vertex whose shape wasn't recognised) must be
// denied, not allowed.
mw := New(providerCfg("gpt-4o"))
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when an allowlist is set")
assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when the provider is restricted")
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_unknown", out.DenyReason.Code, "deny code must be model_unknown")
@@ -122,26 +142,101 @@ func TestAllowlistMissingModelKeyDenies(t *testing.T) {
func TestAllowlistEmptyModelValueDenies(t *testing.T) {
// A present-but-empty model is as undeterminable as an absent one.
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
out, err := mw.Invoke(context.Background(), newInput(
mw := New(providerCfg("gpt-4o"))
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: " "},
))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionDeny, out.Decision, "empty model must be denied when an allowlist is set")
assert.Equal(t, middleware.DecisionDeny, out.Decision, "empty model must be denied when the provider is restricted")
require.NotNil(t, out.DenyReason, "deny reason must be populated")
assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
}
func TestAllowlistEmptyListAllowsMissingModel(t *testing.T) {
// Without an allowlist there is nothing to enforce, so a missing model is
// still allowed — the fail-closed rule only applies when a list is set.
// Without any provider allowlists there is nothing to enforce, so a missing
// model is still allowed — the fail-closed rule only applies when a
// restriction is in effect.
mw := New(Config{})
out, err := mw.Invoke(context.Background(), newInput())
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision, "no allowlist must allow even without a model")
}
func TestUnrestrictedProviderAllowsAnyModel(t *testing.T) {
// The request resolved to otherProvider, which has no allowlist, so its
// traffic must not be caught by testProvider's restriction — the
// cross-provider-leak / false-deny guard.
mw := New(providerCfg("gpt-4o"))
out, err := mw.Invoke(context.Background(), newInputProvider(otherProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision, "an unrestricted provider must not inherit another provider's allowlist")
}
func TestPerProviderAllowlistsAreIsolated(t *testing.T) {
// gpt-4o is allowed only on testProvider; claude-opus-4 only on
// otherProvider. A model allowlisted for one provider must not be usable on
// the other — the fail-closed layer never unions allowlists across providers.
mw := New(Config{ProviderAllowlists: map[string][]string{
testProvider: {"gpt-4o"},
otherProvider: {"claude-opus-4"},
}})
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionDeny, out.Decision, "claude-opus-4 is allowed only on otherProvider, not testProvider")
require.NotNil(t, out.DenyReason)
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "cross-provider model must be blocked, not model_unknown")
}
func TestRestrictionsButNoResolvedProviderFailsClosed(t *testing.T) {
// Restrictions exist for the account but the resolved provider id is absent,
// so the request cannot be scoped to a provider. Fail closed rather than
// wave it through.
mw := New(providerCfg("gpt-4o"))
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.DecisionDeny, out.Decision, "missing resolved provider must fail closed when restrictions exist")
require.NotNil(t, out.DenyReason)
assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
}
func TestEnabledButEmptyAllowlistDeniesEveryModel(t *testing.T) {
// An allowlist-enabled provider with zero models is distinct from an
// unrestricted (absent) provider: it must deny every model.
mw := New(providerCfg())
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionDeny, out.Decision, "an enabled-but-empty allowlist must deny every model")
require.NotNil(t, out.DenyReason)
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must be model_blocked, not model_unknown")
}
func TestFactoryAllEmptyEntriesDenyEveryModel(t *testing.T) {
// All the provider's entries are blank; they collapse to a non-nil empty
// list (deny everything for that provider), not "no restriction".
raw := []byte(`{"provider_allowlists":{"prov-1":[""," "]}}`)
mw, err := Factory{}.New(raw)
require.NoError(t, err)
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionDeny, out.Decision, "all-blank allowlist entries must still restrict the provider")
require.NotNil(t, out.DenyReason)
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must be model_blocked")
}
func TestPromptCaptureDisabledEmitsNoPrompt(t *testing.T) {
mw := New(Config{})
out, err := mw.Invoke(context.Background(), newInput(
@@ -217,8 +312,8 @@ func TestFactoryAcceptsZeroConfigs(t *testing.T) {
func TestFactoryDecodesValidConfig(t *testing.T) {
cfg := Config{
ModelAllowlist: []string{"gpt-4o"},
PromptCapture: PromptCapture{Enabled: true, RedactPii: true},
ProviderAllowlists: map[string][]string{testProvider: {"gpt-4o"}},
PromptCapture: PromptCapture{Enabled: true, RedactPii: true},
}
raw, err := json.Marshal(cfg)
require.NoError(t, err, "marshalling test config must succeed")
@@ -234,15 +329,15 @@ func TestFactoryRejectsMalformedJSON(t *testing.T) {
}
func TestFactoryNormalisesAllowlist(t *testing.T) {
raw := []byte(`{"model_allowlist":[" GPT-4o ","",""," Claude-3 "]}`)
raw := []byte(`{"provider_allowlists":{"prov-1":[" GPT-4o ","",""," Claude-3 "]}}`)
mw, err := Factory{}.New(raw)
require.NoError(t, err)
out, err := mw.Invoke(context.Background(), newInput(
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
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(
out2, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-3"},
))
require.NoError(t, err)

View File

@@ -175,7 +175,7 @@ func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Ou
DenyStatus: 403,
DenyReason: &middleware.DenyReason{
Code: code,
Message: "LLM policy limit exceeded",
Message: denyMessageForCode(code),
},
Metadata: []middleware.KV{
{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},
@@ -184,6 +184,21 @@ func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Ou
}
}
// denyMessageForCode maps a management deny code to a public message.
// Model-allowlist rejections get a model-specific message matching the
// local guardrail; everything else keeps the generic quota wording. The
// message stays generic so it never leaks internal quota detail.
func denyMessageForCode(code string) string {
switch code {
case "llm_policy.model_blocked":
return "model is not in the policy allowlist"
case "llm_policy.model_unknown":
return "request model could not be determined for the policy allowlist"
default:
return "LLM policy limit exceeded"
}
}
// lookupKV returns the value associated with key, or the empty
// string when absent.
func lookupKV(kvs []middleware.KV, key string) string {

View File

@@ -115,6 +115,46 @@ func TestInvoke_DenyConvertsToProxyDeny(t *testing.T) {
assert.NotContains(t, out.DenyReason.Message, "1000", "internal cap numbers must not reach the caller")
}
// TestInvoke_ModelDenyMessages proves a model-allowlist rejection gets a
// model-specific public message rather than the generic quota wording, so a
// blocked or undetermined model reads consistently with the local guardrail.
func TestInvoke_ModelDenyMessages(t *testing.T) {
cases := []struct {
name string
code string
message string
}{
{"blocked", "llm_policy.model_blocked", "model is not in the policy allowlist"},
{"unknown", "llm_policy.model_unknown", "request model could not be determined for the policy allowlist"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
mgmt := &fakeMgmt{
checkResp: &proto.CheckLLMPolicyLimitsResponse{
Decision: "deny",
DenyCode: tc.code,
},
}
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"},
{Key: middleware.KeyLLMModel, Value: "some-model"},
},
})
assert.Equal(t, middleware.DecisionDeny, out.Decision)
require.NotNil(t, out.DenyReason, "deny envelope must carry a reason payload")
assert.Equal(t, tc.code, out.DenyReason.Code, "canonical deny code surfaces to the caller")
assert.Equal(t, tc.message, out.DenyReason.Message,
"model denials must use a model-specific message, matching the local guardrail")
})
}
}
// 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

View File

@@ -6,23 +6,6 @@ import (
"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

View File

@@ -25,10 +25,17 @@ func runParserGuardrail(t *testing.T, url string, body []byte, allowlist []strin
})
require.NoError(t, err, "parser must not error")
guard := llm_guardrail.New(llm_guardrail.Config{ModelAllowlist: allowlist})
const providerID = "prov-under-test"
guard := llm_guardrail.New(llm_guardrail.Config{
ProviderAllowlists: map[string][]string{providerID: allowlist},
})
// The real chain has llm_router stamp the resolved provider id before the
// guardrail runs; the parser doesn't, so add it here so the guardrail can
// scope the allowlist to this provider.
meta := append([]middleware.KV{{Key: middleware.KeyLLMResolvedProviderID, Value: providerID}}, parsed.Metadata...)
out, err := guard.Invoke(context.Background(), &middleware.Input{
Slot: middleware.SlotOnRequest,
Metadata: parsed.Metadata,
Metadata: meta,
})
require.NoError(t, err, "guardrail must not error")
require.NotNil(t, out, "guardrail must return an output")

View File

@@ -8,7 +8,6 @@ package llm_request_parser
import (
"context"
"net/url"
"regexp"
"strconv"
"strings"
"unicode/utf8"
@@ -253,9 +252,7 @@ func parseVertexPath(reqPath string) (vertexRequest, bool) {
if c := strings.LastIndex(rest, ":"); c >= 0 {
model, action = rest[:c], rest[c+1:]
}
if at := strings.Index(model, "@"); at >= 0 {
model = model[:at]
}
model = llm.NormalizeVertexModel(model)
if model == "" {
return vertexRequest{}, false
}
@@ -343,14 +340,6 @@ func trimBedrockNamespace(reqPath string) string {
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:
//
@@ -375,7 +364,7 @@ func parseBedrockPath(reqPath string) (bedrockRequest, bool) {
if decoded, err := url.PathUnescape(rawModel); err == nil {
rawModel = decoded
}
model := normalizeBedrockModel(rawModel)
model := llm.NormalizeBedrockModel(rawModel)
if model == "" {
return bedrockRequest{}, false
}
@@ -389,30 +378,6 @@ func parseBedrockPath(reqPath string) (bedrockRequest, bool) {
}
}
// 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.

View File

@@ -18,10 +18,10 @@ import (
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/test/bufconn"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
mgmtgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
mgmtgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
"github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/proxy/internal/middleware/bodytap"
@@ -134,14 +134,18 @@ func TestReverseProxy_AgentNetworkRequest_FullChain(t *testing.T) {
RedactPii: true,
}))
require.NoError(t, st.SaveAgentNetworkProvider(ctx, &agentNetworkTypes.Provider{
ID: providerID,
AccountID: testAccountID,
ProviderID: "openai_api",
Name: "openai-fullchain-test",
UpstreamURL: upstream.URL, // router rewrites to this
APIKey: "sk-test",
Enabled: true,
Models: []agentNetworkTypes.ProviderModel{{ID: "gpt-5.4"}},
ID: providerID,
AccountID: testAccountID,
ProviderID: "openai_api",
Name: "openai-fullchain-test",
UpstreamURL: upstream.URL, // router rewrites to this
APIKey: "sk-test",
Enabled: true,
// Operator-pinned prices deliberately differ from the catalog's
// gpt-5.4 rates (0.0025/0.015) so the cost assertion below proves
// the per-provider-record price — not the default table — billed
// this request: stored price → synth → wire → cost_meter.
Models: []agentNetworkTypes.ProviderModel{{ID: "gpt-5.4", InputPer1k: 0.004, OutputPer1k: 0.02}},
SessionPrivateKey: "priv",
SessionPublicKey: "pub",
}))
@@ -176,7 +180,7 @@ func TestReverseProxy_AgentNetworkRequest_FullChain(t *testing.T) {
// ---- 5. Wire the middleware framework — same registry the proxy uses
// in production, configured with our bufconn-backed management client.
mwbuiltin.Configure(ctx, t.TempDir(), nil, testLogger, mgmtClient)
mwbuiltin.Configure(ctx, nil, testLogger, mgmtClient)
registry := mwbuiltin.DefaultRegistry()
mwMetrics, err := middleware.NewMetrics(nil)
require.NoError(t, err)
@@ -283,13 +287,23 @@ func TestReverseProxy_AgentNetworkRequest_FullChain(t *testing.T) {
if r.DimensionKind == agentNetworkTypes.DimensionGroup &&
r.DimensionID == adminGroupID &&
r.WindowSeconds == 60 &&
r.TokensInput+r.TokensOutput > 0 {
r.TokensInput+r.TokensOutput > 0 &&
r.CostUSD > 0 {
return true
}
}
return false
}, 5*time.Second, 50*time.Millisecond,
"Admins group consumption row must increment via the response leg — if this fails the proxy's respInput dropped UserGroups again or the parser/recorder wiring is broken")
"Admins group consumption row must increment via the response leg WITH a non-zero cost — a zero cost means the operator's stored price never reached cost_meter (synth → wire → per-record lookup broken)")
// 8a-cost. Exact cost from the OPERATOR's stored price, not the catalog
// default: 12 prompt tokens × 0.004/1k + 40 completion tokens × 0.02/1k
// = 0.000048 + 0.0008 = 0.000848. With catalog rates it would be 0.00063
// — this assertion distinguishes the two, closing the loop on the whole
// dynamic-pricing feature (dashboard save → synth → gRPC wire →
// llm.resolved_provider_id lookup → billing).
assert.Equal(t, "0.000848000", cd.GetMetadata()["cost.usd_total"],
"cost must be computed from the provider record's operator-pinned price")
// 8b. Both the captured prompt and the captured completion are
// redacted — proves the synth threads redact_pii=true into BOTH parser

View File

@@ -246,10 +246,6 @@ type Server struct {
// in processMappings before the receive loop reconnects to resync.
// Zero uses defaultMappingBatchWatchdog.
MappingBatchWatchdog time.Duration
// MiddlewareDataDir is the base directory the middleware system uses to
// resolve file-backed configuration (e.g. the cost_meter pricing table).
// Empty means any middleware that requires a file fails at configure time.
MiddlewareDataDir string
// MiddlewareCaptureBudgetBytes overrides the proxy-wide in-flight capture
// budget passed to middleware.NewManager. Zero or negative values fall
// back to defaultMiddlewareCaptureBudgetBytes (256 MiB).
@@ -2094,7 +2090,7 @@ func (s *Server) initMiddlewareManager(ctx context.Context) error {
return fmt.Errorf("middleware manager requires metrics bundle")
}
otelMeter := s.meter.Meter()
mwbuiltin.Configure(ctx, s.MiddlewareDataDir, otelMeter, s.Logger, s.mgmtClient)
mwbuiltin.Configure(ctx, otelMeter, s.Logger, s.mgmtClient)
mwMetrics, err := middleware.NewMetrics(otelMeter)
if err != nil {