[management, proxy] Management-owned LLM pricing: file-backed defaults + (#6965)

2026-10-24 06:19:08 +02:00 · 2026-07-31 20:52:56 +02:00
parent feecb993f4
commit 0780a806f2
47 changed files with 3196 additions and 1540 deletions
@@ -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")
 }