mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-27 09:09:07 +02:00
[management,proxy] Agent network: per-account LLM gateway (policy, metering, multi-provider) (#6555)
* [agent-network] Shared proto, OpenAPI schema, and generated types * [agent-network] Management: store, manager, synthesizer, policy engine, provider catalog, HTTP/gRPC API Adds the account-scoped agent-network module: provider/policy/budget CRUD and store, the reverse-proxy service synthesizer, policy selection + limit enforcement, the provider catalog (incl. Vertex AI and AWS Bedrock entries), and the management HTTP + proxy gRPC surfaces. * [management] Fix agent-network proxy-peer fan-out on affected-peer recompute The affected-peers resolver loaded only persisted reverse-proxy services, but agent-network services are synthesized on demand and never persisted. As a result the embedded proxy peer was never folded into the affected set when a client's group changed, so the proxy received no network-map update for a newly authorised client and rejected its handshake until a full resync (restart). loadProxyServices now merges the synthesized agent-network services (injected via a registration hook to avoid an import cycle), so proxy peers learn newly authorised clients immediately. * [proxy] Reverse-proxy middleware framework, chain, and request plumbing The per-target middleware chain (slots, dispatcher, mutation gate, metadata merger), body capture, access-log terminal sink, and the proxy wiring that builds + runs chains for synthesized agent-network services. * [proxy] LLM parsers, pricing, and builtin middlewares (OpenAI, Anthropic, Vertex AI, AWS Bedrock) Request/response parsers and SSE/event-stream metering, the embedded pricing table, and the builtin middleware set: request parser, router, policy limit-check/record, cost meter, guardrail, identity inject, response parser. Includes the path-routed providers — Google Vertex AI (keyfile:: service-account OAuth minting) and AWS Bedrock (bearer auth, invoke/converse/streaming, optional /bedrock prefix) — plus the Models allowlist and unmeterable-publisher deny. * [proxy] IPv6 in-place apply and TCP accept-loop hardening on netstack listeners * [agent-network] End-to-end test suite, module docs, and deployment preset * [agent-network] Fix codespell typos and exclude false positives - labelgen word pool: vermillion -> vermilion, racoon -> raccoon. - codespell ignore list: add flate (Go compress/flate package), recordin (a test-local identifier), and unparseable (a valid alternative spelling used consistently across identifiers + a metadata-value constant). * [management] Set LastSeen on injected proxy peer in realstack test (MySQL strict-mode) The injected embedded proxy peer had a PeerStatus with a zero LastSeen, which serializes to '0000-00-00' and is rejected by MySQL in strict mode (SQLite tolerates it). Set LastSeen to a valid time so SaveAccount succeeds on both engines. * [agent-network] Remove e2e shell-script suite from this branch The end-to-end shell scripts under scripts/e2e/ are maintained in a separate testing suite and are not part of this change set. * [agent-network] Polish module docs: remove internal review scaffolding, fix links, verify diagrams Strip PR-review framing, commit references, absolute paths, and stale internal references from the agent-network module docs; fix broken relative links; verify all diagrams against the current architecture. Remove the internal AI-reviewer prompt file. * [management] Refine session expiration handling to support 3-state encoding for SSO deadlines * [agent-network] Relocate agentnetwork package to internals/modules Move management/server/agentnetwork (and its catalog/, labelgen/, types/ subpackages) to management/internals/modules/agentnetwork, alongside the reverse-proxy module, and rewrite all importers. Pure relocation: package names, the synthesizer + affectedpeers registration hook, and store access (shared store.Store) are unchanged, so no import cycle is introduced (affectedpeers still depends only on the agentnetwork/types leaf). * [agent-network] Co-locate HTTP handlers in the module (RegisterEndpoints) Move the agent-network HTTP handlers from server/http/handlers/agentnetwork into the module at internals/modules/agentnetwork/handlers (package handlers) and rename the entrypoint AddEndpoints -> RegisterEndpoints, matching the reverse-proxy module convention. Wiring in http/handler.go updated accordingly.
This commit is contained in:
Maycon Santos
@@ -0,0 +1,65 @@
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestDefaultTable_FirstPartyModelCoverage guards the embedded defaults against
|
||||
// silent drift/gaps: every metered first-party model the management catalog
|
||||
// enumerates must resolve to a price, and a few rates that previously drifted
|
||||
// are pinned to their LiteLLM-validated values. Keep this list in step with the
|
||||
// catalog (management/server/agentnetwork/catalog) when adding models.
|
||||
func TestDefaultTable_FirstPartyModelCoverage(t *testing.T) {
|
||||
tbl := DefaultTable()
|
||||
require.NotNil(t, tbl, "embedded default pricing table must load")
|
||||
|
||||
mustPrice := map[string][]string{
|
||||
// openai parser covers openai_api, azure_openai_api, and mistral_api.
|
||||
"openai": {
|
||||
"gpt-5.5", "gpt-5.5-pro", "gpt-5.4", "gpt-5.4-mini", "gpt-5.4-nano",
|
||||
"gpt-5.3-codex", "gpt-5.3-chat-latest", "o4-mini",
|
||||
"gpt-4.1", "gpt-4.1-mini", "gpt-4.1-nano", "gpt-4o", "gpt-4o-mini",
|
||||
"gpt-4-turbo", "gpt-3.5-turbo", "gpt-35-turbo",
|
||||
"text-embedding-3-large", "text-embedding-3-small",
|
||||
"mistral-large-latest", "mistral-medium-3-5", "codestral-2508",
|
||||
"ministral-8b-latest", "mistral-embed",
|
||||
},
|
||||
"anthropic": {
|
||||
"claude-fable-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6",
|
||||
"claude-opus-4-1", "claude-sonnet-4-6", "claude-sonnet-4-5", "claude-haiku-4-5",
|
||||
},
|
||||
// bedrock keys are the normalized ids the request parser emits.
|
||||
"bedrock": {
|
||||
"anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6",
|
||||
"anthropic.claude-opus-4-1", "anthropic.claude-sonnet-4-6", "anthropic.claude-sonnet-4-5",
|
||||
"anthropic.claude-haiku-4-5", "meta.llama3-3-70b-instruct",
|
||||
"amazon.nova-pro", "amazon.nova-lite", "amazon.nova-micro", "amazon.nova-2-lite",
|
||||
},
|
||||
}
|
||||
for provider, models := range mustPrice {
|
||||
for _, m := range models {
|
||||
_, ok := tbl.Cost(provider, m, 1000, 1000, 0, 0)
|
||||
assert.True(t, ok, "%s/%s must be priced in the embedded defaults", provider, m)
|
||||
}
|
||||
}
|
||||
|
||||
// Pin per-direction rates independently (input-only then output-only) so a
|
||||
// swap or skew of input<->output that preserves the combined total is still
|
||||
// caught — these are rates that previously drifted or are easy to mis-enter.
|
||||
in, ok := tbl.Cost("openai", "gpt-5.4", 1000, 0, 0, 0)
|
||||
require.True(t, ok)
|
||||
assert.InDelta(t, 0.0025, in, 1e-9, "gpt-5.4 input = 0.0025 per 1k")
|
||||
out, ok := tbl.Cost("openai", "gpt-5.4", 0, 1000, 0, 0)
|
||||
require.True(t, ok)
|
||||
assert.InDelta(t, 0.015, out, 1e-9, "gpt-5.4 output = 0.015 per 1k")
|
||||
|
||||
in, ok = tbl.Cost("bedrock", "anthropic.claude-sonnet-4-5", 1000, 0, 0, 0)
|
||||
require.True(t, ok)
|
||||
assert.InDelta(t, 0.003, in, 1e-9, "bedrock sonnet-4-5 input = 0.003 per 1k")
|
||||
out, ok = tbl.Cost("bedrock", "anthropic.claude-sonnet-4-5", 0, 1000, 0, 0)
|
||||
require.True(t, ok)
|
||||
assert.InDelta(t, 0.015, out, 1e-9, "bedrock sonnet-4-5 output = 0.015 per 1k")
|
||||
}
|
||||
@@ -0,0 +1,264 @@
|
||||
# Embedded default pricing for llm_observability. Compiled into the proxy
|
||||
# binary via go:embed in pricing.go; cost annotation works out of the box
|
||||
# without any operator action.
|
||||
#
|
||||
# Operators override entries by dropping a pricing.yaml into --plugin-data-dir
|
||||
# (or whichever basename is given via params.pricing_path). The override file
|
||||
# only needs entries the operator wants to change; missing entries fall
|
||||
# through to these defaults.
|
||||
#
|
||||
# Values are USD per 1_000 tokens. Public list prices drift; ship a fresh
|
||||
# binary or override individual entries via the override file as needed.
|
||||
#
|
||||
# Optional cache fields:
|
||||
# cached_input_per_1k OpenAI: rate for prompt_tokens_details.cached_tokens
|
||||
# (a SUBSET of prompt_tokens). Typically 0.5x input.
|
||||
# Absent → cached portion bills at input_per_1k.
|
||||
# cache_read_per_1k Anthropic: rate for cache_read_input_tokens
|
||||
# (ADDITIVE to input_tokens). Typically 0.1x input.
|
||||
# Absent → cache reads bill at input_per_1k.
|
||||
# cache_creation_per_1k Anthropic: rate for cache_creation_input_tokens
|
||||
# (ADDITIVE to input_tokens). Typically 1.25x input.
|
||||
# Absent → cache writes bill at input_per_1k.
|
||||
|
||||
openai:
|
||||
# OpenAI + OpenAI-compatible providers (openai_api, azure_openai_api,
|
||||
# mistral_api, and the openai-parser gateways) all emit llm.provider="openai",
|
||||
# so their models are priced here. Kept in sync with the management catalog;
|
||||
# rates cross-checked against LiteLLM model_prices_and_context_window.json.
|
||||
|
||||
# GPT-5.x family — cache reads 10% of input (0.1x).
|
||||
gpt-5.5:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.03
|
||||
cached_input_per_1k: 0.0005
|
||||
gpt-5.5-pro:
|
||||
input_per_1k: 0.03
|
||||
output_per_1k: 0.18
|
||||
cached_input_per_1k: 0.003
|
||||
gpt-5.4:
|
||||
input_per_1k: 0.0025
|
||||
output_per_1k: 0.015
|
||||
cached_input_per_1k: 0.00025
|
||||
gpt-5.4-pro:
|
||||
input_per_1k: 0.03
|
||||
output_per_1k: 0.18
|
||||
cached_input_per_1k: 0.003
|
||||
gpt-5.4-mini:
|
||||
input_per_1k: 0.00075
|
||||
output_per_1k: 0.0045
|
||||
cached_input_per_1k: 0.000075
|
||||
gpt-5.4-nano:
|
||||
input_per_1k: 0.0002
|
||||
output_per_1k: 0.00125
|
||||
cached_input_per_1k: 0.00002
|
||||
gpt-5.3-codex:
|
||||
input_per_1k: 0.00175
|
||||
output_per_1k: 0.014
|
||||
cached_input_per_1k: 0.000175
|
||||
gpt-5.3-chat-latest:
|
||||
input_per_1k: 0.00175
|
||||
output_per_1k: 0.014
|
||||
cached_input_per_1k: 0.000175
|
||||
# GPT-5 (2025) family — kept for gateway requests using the unsuffixed ids.
|
||||
gpt-5:
|
||||
input_per_1k: 0.00125
|
||||
output_per_1k: 0.01
|
||||
cached_input_per_1k: 0.000125
|
||||
gpt-5-mini:
|
||||
input_per_1k: 0.00025
|
||||
output_per_1k: 0.002
|
||||
cached_input_per_1k: 0.000025
|
||||
gpt-5-nano:
|
||||
input_per_1k: 0.00005
|
||||
output_per_1k: 0.0004
|
||||
cached_input_per_1k: 0.000005
|
||||
o4-mini:
|
||||
input_per_1k: 0.0011
|
||||
output_per_1k: 0.0044
|
||||
cached_input_per_1k: 0.000275
|
||||
# GPT-4.1 family — cache reads 25% of input.
|
||||
gpt-4.1:
|
||||
input_per_1k: 0.002
|
||||
output_per_1k: 0.008
|
||||
cached_input_per_1k: 0.0005
|
||||
gpt-4.1-mini:
|
||||
input_per_1k: 0.0004
|
||||
output_per_1k: 0.0016
|
||||
cached_input_per_1k: 0.0001
|
||||
gpt-4.1-nano:
|
||||
input_per_1k: 0.0001
|
||||
output_per_1k: 0.0004
|
||||
cached_input_per_1k: 0.000025
|
||||
# GPT-4o family — cache reads 50% of input (0.5x).
|
||||
gpt-4o:
|
||||
input_per_1k: 0.0025
|
||||
output_per_1k: 0.01
|
||||
cached_input_per_1k: 0.00125
|
||||
gpt-4o-mini:
|
||||
input_per_1k: 0.00015
|
||||
output_per_1k: 0.0006
|
||||
cached_input_per_1k: 0.000075
|
||||
# Older GPT — no prompt caching.
|
||||
gpt-4-turbo:
|
||||
input_per_1k: 0.01
|
||||
output_per_1k: 0.03
|
||||
gpt-3.5-turbo:
|
||||
input_per_1k: 0.0005
|
||||
output_per_1k: 0.0015
|
||||
gpt-35-turbo: # Azure deployment alias of gpt-3.5-turbo
|
||||
input_per_1k: 0.0005
|
||||
output_per_1k: 0.0015
|
||||
# Embeddings — no caching, no output tokens.
|
||||
text-embedding-3-large:
|
||||
input_per_1k: 0.00013
|
||||
output_per_1k: 0
|
||||
text-embedding-3-small:
|
||||
input_per_1k: 0.00002
|
||||
output_per_1k: 0
|
||||
|
||||
# Mistral (mistral_api) — routed via the openai parser; no prompt caching.
|
||||
mistral-large-latest:
|
||||
input_per_1k: 0.0005
|
||||
output_per_1k: 0.0015
|
||||
mistral-medium-latest:
|
||||
input_per_1k: 0.0004
|
||||
output_per_1k: 0.002
|
||||
mistral-medium-3-5:
|
||||
input_per_1k: 0.0015
|
||||
output_per_1k: 0.0075
|
||||
mistral-small-latest:
|
||||
input_per_1k: 0.00006
|
||||
output_per_1k: 0.00018
|
||||
magistral-medium-latest:
|
||||
input_per_1k: 0.002
|
||||
output_per_1k: 0.005
|
||||
magistral-small-latest:
|
||||
input_per_1k: 0.0005
|
||||
output_per_1k: 0.0015
|
||||
devstral-medium-latest:
|
||||
input_per_1k: 0.0004
|
||||
output_per_1k: 0.002
|
||||
devstral-small-latest:
|
||||
input_per_1k: 0.0001
|
||||
output_per_1k: 0.0003
|
||||
codestral-2508:
|
||||
input_per_1k: 0.0003
|
||||
output_per_1k: 0.0009
|
||||
codestral-latest:
|
||||
input_per_1k: 0.001
|
||||
output_per_1k: 0.003
|
||||
ministral-3-14b-2512:
|
||||
input_per_1k: 0.0002
|
||||
output_per_1k: 0.0002
|
||||
ministral-8b-latest:
|
||||
input_per_1k: 0.00015
|
||||
output_per_1k: 0.00015
|
||||
ministral-3-3b-2512:
|
||||
input_per_1k: 0.0001
|
||||
output_per_1k: 0.0001
|
||||
mistral-embed:
|
||||
input_per_1k: 0.0001
|
||||
output_per_1k: 0
|
||||
|
||||
anthropic:
|
||||
# Claude 4.x family — cache reads ≈10% of input, cache writes ≈125% of input.
|
||||
# Pricing source: Anthropic's current published rates per million tokens,
|
||||
# divided by 1000 for the per-1k figures stored here.
|
||||
claude-fable-5:
|
||||
input_per_1k: 0.010
|
||||
output_per_1k: 0.050
|
||||
cache_read_per_1k: 0.001
|
||||
cache_creation_per_1k: 0.0125
|
||||
claude-opus-4-8:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
claude-opus-4-7:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
claude-opus-4-6:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
claude-opus-4-1:
|
||||
input_per_1k: 0.015
|
||||
output_per_1k: 0.075
|
||||
cache_read_per_1k: 0.0015
|
||||
cache_creation_per_1k: 0.01875
|
||||
claude-sonnet-4-6:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
claude-sonnet-4-5:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
claude-haiku-4-5:
|
||||
input_per_1k: 0.001
|
||||
output_per_1k: 0.005
|
||||
cache_read_per_1k: 0.0001
|
||||
cache_creation_per_1k: 0.00125
|
||||
|
||||
bedrock:
|
||||
# AWS Bedrock model ids, normalised by the request parser (cross-region
|
||||
# inference-profile prefix + version/throughput suffix stripped), e.g.
|
||||
# eu.anthropic.claude-sonnet-4-5-20250929-v1:0 -> anthropic.claude-sonnet-4-5.
|
||||
# Anthropic-on-Bedrock keeps the additive cache buckets (read ≈0.1x input,
|
||||
# write ≈1.25x input); Nova / Llama report no cache, so cost is input+output.
|
||||
anthropic.claude-opus-4-8:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
anthropic.claude-opus-4-7:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
anthropic.claude-opus-4-6:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
anthropic.claude-opus-4-1:
|
||||
input_per_1k: 0.015
|
||||
output_per_1k: 0.075
|
||||
cache_read_per_1k: 0.0015
|
||||
cache_creation_per_1k: 0.01875
|
||||
anthropic.claude-sonnet-4-6:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
anthropic.claude-sonnet-4-5:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
anthropic.claude-haiku-4-5:
|
||||
input_per_1k: 0.001
|
||||
output_per_1k: 0.005
|
||||
cache_read_per_1k: 0.0001
|
||||
cache_creation_per_1k: 0.00125
|
||||
meta.llama3-3-70b-instruct:
|
||||
input_per_1k: 0.00072
|
||||
output_per_1k: 0.00072
|
||||
amazon.nova-2-lite:
|
||||
input_per_1k: 0.0003
|
||||
output_per_1k: 0.0025
|
||||
amazon.nova-pro:
|
||||
input_per_1k: 0.0008
|
||||
output_per_1k: 0.0032
|
||||
amazon.nova-lite:
|
||||
input_per_1k: 0.00006
|
||||
output_per_1k: 0.00024
|
||||
amazon.nova-micro:
|
||||
input_per_1k: 0.000035
|
||||
output_per_1k: 0.00014
|
||||
@@ -0,0 +1,449 @@
|
||||
// Package pricing implements the embedded-default + override pricing table
|
||||
// shared by middleware that converts LLM token usage into a USD cost
|
||||
// estimate. The table is hot-reloadable from a basename under the proxy
|
||||
// data directory; missing override files keep the embedded defaults so
|
||||
// cost annotation works without operator action.
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
_ "embed"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/fs"
|
||||
"math"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/metric"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
//go:embed defaults_pricing.yaml
|
||||
var defaultPricingYAML []byte
|
||||
|
||||
var (
|
||||
defaultTableOnce sync.Once
|
||||
defaultTablePtr *Table
|
||||
)
|
||||
|
||||
// DefaultTable returns the pricing table embedded in the binary. The result
|
||||
// is parsed once and shared; callers must not mutate the returned value.
|
||||
// Cost annotation works without any operator action because every loader
|
||||
// starts with this table.
|
||||
func DefaultTable() *Table {
|
||||
defaultTableOnce.Do(func() {
|
||||
t, err := parsePricingBytes(defaultPricingYAML)
|
||||
if err != nil {
|
||||
panic(fmt.Sprintf("llmobs: embedded default pricing failed to parse: %v", err))
|
||||
}
|
||||
defaultTablePtr = t
|
||||
})
|
||||
return defaultTablePtr
|
||||
}
|
||||
|
||||
// mergeOver returns a new Table containing every entry from base, with any
|
||||
// matching entry from overlay replacing the base value. Either argument may
|
||||
// be nil. Result is a fresh allocation so callers can mutate / Store safely.
|
||||
func mergeOver(base, overlay *Table) *Table {
|
||||
if overlay == nil || len(overlay.entries) == 0 {
|
||||
return base
|
||||
}
|
||||
if base == nil || len(base.entries) == 0 {
|
||||
return overlay
|
||||
}
|
||||
out := make(map[string]map[string]Entry, len(base.entries))
|
||||
for provider, models := range base.entries {
|
||||
inner := make(map[string]Entry, len(models))
|
||||
for model, e := range models {
|
||||
inner[model] = e
|
||||
}
|
||||
out[provider] = inner
|
||||
}
|
||||
for provider, models := range overlay.entries {
|
||||
inner, ok := out[provider]
|
||||
if !ok {
|
||||
inner = make(map[string]Entry, len(models))
|
||||
out[provider] = inner
|
||||
}
|
||||
for model, e := range models {
|
||||
inner[model] = e
|
||||
}
|
||||
}
|
||||
return &Table{entries: out}
|
||||
}
|
||||
|
||||
// Entry is a single model's input and output pricing, expressed in USD per
|
||||
// 1000 tokens.
|
||||
//
|
||||
// CachedInputPer1K applies to OpenAI's cached prompt tokens, which are a
|
||||
// subset of input_tokens — when set, the cached portion is billed at this
|
||||
// rate and the non-cached remainder at InputPer1K. Zero means "no discount
|
||||
// configured", and cached tokens are billed at InputPer1K (matches current
|
||||
// behaviour where cached counts weren't extracted at all).
|
||||
//
|
||||
// CacheReadPer1K and CacheCreationPer1K apply to Anthropic's two prompt-
|
||||
// cache fields, which are additive to input_tokens: cache_read is the
|
||||
// cheaper read-from-cache rate, cache_creation is the more expensive
|
||||
// write-to-cache rate. Zero means "no rate configured" and the
|
||||
// corresponding token bucket is billed at InputPer1K. This is more
|
||||
// accurate than today's behaviour, where Anthropic's cache tokens are
|
||||
// ignored and not charged at all.
|
||||
type Entry struct {
|
||||
InputPer1K float64
|
||||
OutputPer1K float64
|
||||
CachedInputPer1K float64
|
||||
CacheReadPer1K float64
|
||||
CacheCreationPer1K float64
|
||||
}
|
||||
|
||||
// Table is a provider-to-model pricing lookup. Instances are immutable once
|
||||
// built and are swapped atomically by Loader.
|
||||
type Table struct {
|
||||
entries map[string]map[string]Entry
|
||||
}
|
||||
|
||||
// Cost returns the estimated USD cost for the given token counts. ok is
|
||||
// false when the provider or model is not present in the table; the caller
|
||||
// can still emit token metrics with a model=unknown label.
|
||||
//
|
||||
// Provider-shape semantics for cached / cache-creation counts:
|
||||
//
|
||||
// - OpenAI: cachedInput is a SUBSET of inTokens. The cached portion is
|
||||
// billed at CachedInputPer1K (or InputPer1K when no override), and the
|
||||
// non-cached remainder of inTokens at InputPer1K. cacheCreation is
|
||||
// ignored (OpenAI has no analogue).
|
||||
// - Anthropic: cachedInput (cache_read) and cacheCreation are ADDITIVE to
|
||||
// inTokens. The three buckets are billed at CacheReadPer1K,
|
||||
// CacheCreationPer1K, and InputPer1K respectively, each falling back
|
||||
// to InputPer1K when the corresponding rate is zero.
|
||||
// - Other providers: cached and cacheCreation are ignored; cost is
|
||||
// inTokens*InputPer1K + outTokens*OutputPer1K.
|
||||
func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (float64, bool) {
|
||||
// Clamp negatives to zero before any pricing math so a malformed
|
||||
// upstream count can never produce a negative cost.
|
||||
if inTokens < 0 {
|
||||
inTokens = 0
|
||||
}
|
||||
if outTokens < 0 {
|
||||
outTokens = 0
|
||||
}
|
||||
if cachedInput < 0 {
|
||||
cachedInput = 0
|
||||
}
|
||||
if cacheCreation < 0 {
|
||||
cacheCreation = 0
|
||||
}
|
||||
if t == nil {
|
||||
return 0, false
|
||||
}
|
||||
byModel, ok := t.entries[provider]
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
entry, ok := byModel[model]
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
output := (float64(outTokens) / 1000.0) * entry.OutputPer1K
|
||||
switch provider {
|
||||
case "openai":
|
||||
// cachedInput is a subset of inTokens; clamp so a malformed
|
||||
// upstream (cached > total) can't produce a negative remainder.
|
||||
clamped := cachedInput
|
||||
if clamped > inTokens {
|
||||
clamped = inTokens
|
||||
}
|
||||
cachedRate := entry.CachedInputPer1K
|
||||
if cachedRate <= 0 {
|
||||
cachedRate = entry.InputPer1K
|
||||
}
|
||||
nonCached := float64(inTokens-clamped) / 1000.0 * entry.InputPer1K
|
||||
cached := float64(clamped) / 1000.0 * cachedRate
|
||||
return nonCached + cached + output, true
|
||||
case "anthropic", "bedrock":
|
||||
// Bedrock-Anthropic returns the same additive cache buckets as
|
||||
// first-party Anthropic; non-Anthropic Bedrock models simply report
|
||||
// zero cache tokens, so this formula degrades to input + output.
|
||||
readRate := entry.CacheReadPer1K
|
||||
if readRate <= 0 {
|
||||
readRate = entry.InputPer1K
|
||||
}
|
||||
createRate := entry.CacheCreationPer1K
|
||||
if createRate <= 0 {
|
||||
createRate = entry.InputPer1K
|
||||
}
|
||||
input := float64(inTokens) / 1000.0 * entry.InputPer1K
|
||||
read := float64(cachedInput) / 1000.0 * readRate
|
||||
create := float64(cacheCreation) / 1000.0 * createRate
|
||||
return input + read + create + output, true
|
||||
default:
|
||||
input := float64(inTokens) / 1000.0 * entry.InputPer1K
|
||||
return input + output, true
|
||||
}
|
||||
}
|
||||
|
||||
// Has reports whether the provider/model pair is present in the table.
|
||||
func (t *Table) Has(provider, model string) bool {
|
||||
if t == nil {
|
||||
return false
|
||||
}
|
||||
byModel, ok := t.entries[provider]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
_, ok = byModel[model]
|
||||
return ok
|
||||
}
|
||||
|
||||
// pricingFile mirrors the on-disk YAML schema. Keys are provider names; the
|
||||
// nested map keys are model names.
|
||||
type pricingFile map[string]map[string]struct {
|
||||
InputPer1K float64 `yaml:"input_per_1k"`
|
||||
OutputPer1K float64 `yaml:"output_per_1k"`
|
||||
CachedInputPer1K float64 `yaml:"cached_input_per_1k"`
|
||||
CacheReadPer1K float64 `yaml:"cache_read_per_1k"`
|
||||
CacheCreationPer1K float64 `yaml:"cache_creation_per_1k"`
|
||||
}
|
||||
|
||||
const (
|
||||
// ReloadInterval is the mtime-poll cadence for the background reloader.
|
||||
ReloadInterval = 30 * time.Second
|
||||
|
||||
// errorBackoff bounds how often the loader logs a repeated parse error.
|
||||
errorBackoff = 5 * time.Minute
|
||||
)
|
||||
|
||||
var basenameRegex = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)
|
||||
|
||||
// Loader is a confined, hot-reloadable pricing table reader. Construction
|
||||
// must succeed against the target file; subsequent reload failures keep the
|
||||
// previously-loaded table so callers never observe a blank price list.
|
||||
type Loader struct {
|
||||
baseDir string
|
||||
fullPath string
|
||||
pluginID string
|
||||
table atomic.Pointer[Table]
|
||||
mtime atomic.Int64
|
||||
failures metric.Int64Counter
|
||||
interval time.Duration
|
||||
}
|
||||
|
||||
// NewLoader returns a pricing loader that overlays an optional file-based
|
||||
// table on top of the embedded defaults. Missing override file, baseDir, or
|
||||
// relPath is not an error: the loader keeps the embedded defaults so cost
|
||||
// metadata is still emitted for known models.
|
||||
//
|
||||
// Errors:
|
||||
// - bad basename, traversal segment, or absolute relPath are rejected so a
|
||||
// misconfigured target surfaces immediately.
|
||||
// - permission errors and YAML parse errors keep the defaults but log a
|
||||
// warning; cost annotation does not silently break.
|
||||
//
|
||||
// failures is optional; pass nil in tests that do not care about
|
||||
// reload-failure telemetry.
|
||||
func NewLoader(baseDir, relPath, pluginID string, failures metric.Int64Counter) (*Loader, error) {
|
||||
defaults := DefaultTable()
|
||||
l := &Loader{
|
||||
baseDir: baseDir,
|
||||
pluginID: pluginID,
|
||||
failures: failures,
|
||||
}
|
||||
l.table.Store(defaults)
|
||||
|
||||
if strings.TrimSpace(baseDir) == "" || strings.TrimSpace(relPath) == "" {
|
||||
return l, nil
|
||||
}
|
||||
|
||||
full, err := resolveMiddlewareDataPath(baseDir, relPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
l.fullPath = full
|
||||
|
||||
overlay, mtime, err := loadPricing(full)
|
||||
if err != nil {
|
||||
if errors.Is(err, fs.ErrNotExist) {
|
||||
// Override file is optional. Defaults already stored.
|
||||
return l, nil
|
||||
}
|
||||
// Symlink rejection, oversize file, parse failure, permission errors
|
||||
// — surface so a misconfigured operator sees the problem instead of
|
||||
// silently running with stale defaults.
|
||||
return nil, fmt.Errorf("load pricing %s: %w", full, err)
|
||||
}
|
||||
l.table.Store(mergeOver(defaults, overlay))
|
||||
l.mtime.Store(mtime.UnixNano())
|
||||
return l, nil
|
||||
}
|
||||
|
||||
// Get returns the current pricing table. The returned pointer is immutable;
|
||||
// callers must not mutate its contents.
|
||||
func (l *Loader) Get() *Table {
|
||||
if l == nil {
|
||||
return nil
|
||||
}
|
||||
return l.table.Load()
|
||||
}
|
||||
|
||||
// WatchesFile reports whether this loader is bound to an override file on
|
||||
// disk. False for defaults-only loaders (no operator override given).
|
||||
// Callers use this to decide whether to spawn the mtime-poll goroutine.
|
||||
func (l *Loader) WatchesFile() bool {
|
||||
if l == nil {
|
||||
return false
|
||||
}
|
||||
return l.fullPath != ""
|
||||
}
|
||||
|
||||
// SetReloadInterval overrides the mtime-poll cadence used by Reload. Calls
|
||||
// after Reload has started have no effect on the running loop. Intended for
|
||||
// tests; production code uses the default ReloadInterval.
|
||||
func (l *Loader) SetReloadInterval(d time.Duration) {
|
||||
if l == nil || d <= 0 {
|
||||
return
|
||||
}
|
||||
l.interval = d
|
||||
}
|
||||
|
||||
// Reload runs a polling loop that checks the pricing file mtime every
|
||||
// ReloadInterval (or the value passed to SetReloadInterval). Returns when
|
||||
// ctx is cancelled.
|
||||
func (l *Loader) Reload(ctx context.Context) {
|
||||
if l == nil {
|
||||
return
|
||||
}
|
||||
interval := l.interval
|
||||
if interval <= 0 {
|
||||
interval = ReloadInterval
|
||||
}
|
||||
t := time.NewTicker(interval)
|
||||
defer t.Stop()
|
||||
|
||||
var lastErrAt time.Time
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
if err := l.reload(); err != nil {
|
||||
if l.failures != nil {
|
||||
l.failures.Add(ctx, 1, metric.WithAttributes(
|
||||
attribute.String("plugin", l.pluginID),
|
||||
))
|
||||
}
|
||||
now := time.Now()
|
||||
if now.Sub(lastErrAt) >= errorBackoff {
|
||||
log.Warnf("llmobs: pricing reload failed for %s: %v", l.fullPath, err)
|
||||
lastErrAt = now
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// reload performs a single-shot mtime check and reload. The reloaded
|
||||
// override file is merged on top of the embedded defaults; missing override
|
||||
// (e.g. operator deleted the file) is not an error and reverts to defaults.
|
||||
func (l *Loader) reload() error {
|
||||
if l.fullPath == "" {
|
||||
// Defaults-only loader; nothing on disk to reload.
|
||||
return nil
|
||||
}
|
||||
mtime, err := statMtime(l.fullPath)
|
||||
if err != nil {
|
||||
if errors.Is(err, fs.ErrNotExist) {
|
||||
// File was removed since startup. Drop back to defaults and
|
||||
// reset mtime so a future re-creation triggers a reload.
|
||||
l.table.Store(DefaultTable())
|
||||
l.mtime.Store(0)
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
if mtime.UnixNano() == l.mtime.Load() {
|
||||
return nil
|
||||
}
|
||||
|
||||
overlay, newMtime, err := loadPricing(l.fullPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
l.table.Store(mergeOver(DefaultTable(), overlay))
|
||||
l.mtime.Store(newMtime.UnixNano())
|
||||
return nil
|
||||
}
|
||||
|
||||
// resolveMiddlewareDataPath validates relPath is a safe basename and resolves
|
||||
// it under baseDir. An additional cleaned-prefix check guards against
|
||||
// CVE-style edge cases where Join is used with trailing path segments.
|
||||
func resolveMiddlewareDataPath(baseDir, relPath string) (string, error) {
|
||||
if strings.TrimSpace(baseDir) == "" {
|
||||
return "", errors.New("middleware-data-dir is not configured")
|
||||
}
|
||||
if relPath == "" {
|
||||
return "", errors.New("pricing path is empty")
|
||||
}
|
||||
if !basenameRegex.MatchString(relPath) {
|
||||
return "", fmt.Errorf("pricing path %q is not a safe basename", relPath)
|
||||
}
|
||||
if filepath.IsAbs(relPath) {
|
||||
return "", fmt.Errorf("pricing path %q must be a basename, not absolute", relPath)
|
||||
}
|
||||
|
||||
cleanBase, err := filepath.Abs(filepath.Clean(baseDir))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("resolve middleware-data-dir: %w", err)
|
||||
}
|
||||
full := filepath.Join(cleanBase, relPath)
|
||||
cleanedFull := filepath.Clean(full)
|
||||
if !strings.HasPrefix(cleanedFull, cleanBase+string(filepath.Separator)) && cleanedFull != cleanBase {
|
||||
return "", fmt.Errorf("pricing path %q escapes middleware-data-dir", relPath)
|
||||
}
|
||||
return cleanedFull, nil
|
||||
}
|
||||
|
||||
func parsePricingBytes(data []byte) (*Table, error) {
|
||||
dec := yaml.NewDecoder(bytes.NewReader(data))
|
||||
dec.KnownFields(true)
|
||||
|
||||
var raw pricingFile
|
||||
if err := dec.Decode(&raw); err != nil && !errors.Is(err, io.EOF) {
|
||||
return nil, fmt.Errorf("decode pricing yaml: %w", err)
|
||||
}
|
||||
|
||||
out := make(map[string]map[string]Entry, len(raw))
|
||||
for provider, models := range raw {
|
||||
inner := make(map[string]Entry, len(models))
|
||||
for model, entry := range models {
|
||||
for field, v := range map[string]float64{
|
||||
"input_per_1k": entry.InputPer1K,
|
||||
"output_per_1k": entry.OutputPer1K,
|
||||
"cached_input_per_1k": entry.CachedInputPer1K,
|
||||
"cache_read_per_1k": entry.CacheReadPer1K,
|
||||
"cache_creation_per_1k": entry.CacheCreationPer1K,
|
||||
} {
|
||||
if v < 0 || math.IsNaN(v) || math.IsInf(v, 0) {
|
||||
return nil, fmt.Errorf("pricing %s/%s: %s must be a finite, non-negative rate, got %v", provider, model, field, v)
|
||||
}
|
||||
}
|
||||
inner[model] = Entry{
|
||||
InputPer1K: entry.InputPer1K,
|
||||
OutputPer1K: entry.OutputPer1K,
|
||||
CachedInputPer1K: entry.CachedInputPer1K,
|
||||
CacheReadPer1K: entry.CacheReadPer1K,
|
||||
CacheCreationPer1K: entry.CacheCreationPer1K,
|
||||
}
|
||||
}
|
||||
out[provider] = inner
|
||||
}
|
||||
return &Table{entries: out}, nil
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
//go:build !unix
|
||||
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// loadPricing is unavailable on non-Unix platforms because O_NOFOLLOW and
|
||||
// fstat-from-FD are required to honour the spec's symlink-safety rules. The
|
||||
// proxy is only deployed on Linux today; a Windows port would need an
|
||||
// equivalent path-as-handle implementation.
|
||||
func loadPricing(path string) (*Table, time.Time, error) {
|
||||
return nil, time.Time{}, fmt.Errorf("llmobs pricing loader is not supported on this platform: %s", path)
|
||||
}
|
||||
|
||||
func statMtime(path string) (time.Time, error) {
|
||||
return time.Time{}, fmt.Errorf("llmobs pricing loader is not supported on this platform: %s", path)
|
||||
}
|
||||
@@ -0,0 +1,432 @@
|
||||
//go:build unix
|
||||
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func copyFixture(t *testing.T, src, dst string) {
|
||||
t.Helper()
|
||||
data, err := os.ReadFile(src)
|
||||
require.NoError(t, err, "read source fixture")
|
||||
require.NoError(t, os.WriteFile(dst, data, 0o600), "write target fixture")
|
||||
}
|
||||
|
||||
func TestNewLoader_HappyPath(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
|
||||
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err, "NewLoader must succeed with a valid fixture")
|
||||
table := l.Get()
|
||||
require.NotNil(t, table, "table populated after load")
|
||||
|
||||
cost, ok := table.Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
|
||||
require.True(t, ok, "known provider/model resolves")
|
||||
assert.InDelta(t, 0.00075, cost, 1e-9, "cost = 0.00015 + 0.0006 per 1k tokens")
|
||||
|
||||
cost, ok = table.Cost("openai", "gpt-4o", 2000, 1000, 0, 0)
|
||||
require.True(t, ok, "second known model resolves")
|
||||
assert.InDelta(t, 0.015, cost, 1e-9, "cost for gpt-4o: 2*0.0025 + 1*0.01")
|
||||
|
||||
cost, ok = table.Cost("anthropic", "claude-sonnet-4-5", 1000, 1000, 0, 0)
|
||||
require.True(t, ok, "anthropic model resolves")
|
||||
assert.InDelta(t, 0.018, cost, 1e-9, "cost for claude-sonnet-4-5: 0.003 + 0.015")
|
||||
}
|
||||
|
||||
// TestCost_OpenAICachedSubsetDiscount proves OpenAI's cached input
|
||||
// tokens are billed at the configured cached_input_per_1k rate while
|
||||
// the non-cached remainder of input_tokens is billed at the regular
|
||||
// rate. Critical because OpenAI returns cached_tokens as a SUBSET of
|
||||
// prompt_tokens — naïvely charging the cached count on top of
|
||||
// prompt_tokens would double-bill that portion.
|
||||
func TestCost_OpenAICachedSubsetDiscount(t *testing.T) {
|
||||
tbl := &Table{entries: map[string]map[string]Entry{
|
||||
"openai": {"gpt-4o": {
|
||||
InputPer1K: 0.0025, // 0.0025 USD per 1k input tokens
|
||||
OutputPer1K: 0.01,
|
||||
CachedInputPer1K: 0.00125, // 0.5x discount on cached
|
||||
}},
|
||||
}}
|
||||
// 1000 prompt tokens, 750 of which were cached. 250 non-cached
|
||||
// at regular rate, 750 cached at the discount rate, 500 output.
|
||||
cost, ok := tbl.Cost("openai", "gpt-4o", 1000, 500, 750, 0)
|
||||
require.True(t, ok, "known model resolves")
|
||||
want := (250.0/1000.0)*0.0025 + (750.0/1000.0)*0.00125 + (500.0/1000.0)*0.01
|
||||
assert.InDelta(t, want, cost, 1e-12,
|
||||
"cached subset must bill at the discount rate; non-cached remainder at regular rate")
|
||||
}
|
||||
|
||||
// TestCost_OpenAICachedFallsBackToInputRate covers the operator
|
||||
// opt-in contract: when CachedInputPer1K is unset (zero), cached
|
||||
// tokens bill at the regular input rate. This matches today's
|
||||
// behaviour (cached counts weren't extracted at all so they
|
||||
// implicitly billed at the input rate via prompt_tokens).
|
||||
func TestCost_OpenAICachedFallsBackToInputRate(t *testing.T) {
|
||||
tbl := &Table{entries: map[string]map[string]Entry{
|
||||
"openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01}},
|
||||
}}
|
||||
cost, ok := tbl.Cost("openai", "gpt-4o", 1000, 500, 750, 0)
|
||||
require.True(t, ok)
|
||||
want := 0.0025 + (500.0/1000.0)*0.01
|
||||
assert.InDelta(t, want, cost, 1e-12,
|
||||
"absent cached_input_per_1k rate must fall back to input_per_1k — same as pre-feature behaviour")
|
||||
}
|
||||
|
||||
// TestCost_OpenAIClampsCachedToInputCount is the defensive guard
|
||||
// against malformed upstream responses that report cached_tokens >
|
||||
// prompt_tokens. We clamp so the formula never produces a negative
|
||||
// "non-cached remainder" multiplied by the input rate.
|
||||
func TestCost_OpenAIClampsCachedToInputCount(t *testing.T) {
|
||||
tbl := &Table{entries: map[string]map[string]Entry{
|
||||
"openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01, CachedInputPer1K: 0.00125}},
|
||||
}}
|
||||
cost, ok := tbl.Cost("openai", "gpt-4o", 100, 0, 9999, 0)
|
||||
require.True(t, ok)
|
||||
// All 100 cached, 0 non-cached. Output is 0.
|
||||
want := (100.0 / 1000.0) * 0.00125
|
||||
assert.InDelta(t, want, cost, 1e-12,
|
||||
"cached count > input count must clamp to input — never bill negative non-cached tokens")
|
||||
}
|
||||
|
||||
// TestCost_AnthropicCacheReadAndCreationAreAdditive proves the
|
||||
// Anthropic shape: cache_read and cache_creation tokens are
|
||||
// ADDITIVE to input_tokens (not subset), each billed at its own
|
||||
// configured rate. The two rates pull in opposite directions —
|
||||
// cache_read is the cheaper read-from-cache rate (≈0.1× input),
|
||||
// cache_creation is the more expensive write-to-cache rate
|
||||
// (≈1.25× input).
|
||||
func TestCost_AnthropicCacheReadAndCreationAreAdditive(t *testing.T) {
|
||||
tbl := &Table{entries: map[string]map[string]Entry{
|
||||
"anthropic": {"claude-sonnet": {
|
||||
InputPer1K: 0.003,
|
||||
OutputPer1K: 0.015,
|
||||
CacheReadPer1K: 0.0003, // 0.1x of input
|
||||
CacheCreationPer1K: 0.00375, // 1.25x of input
|
||||
}},
|
||||
}}
|
||||
// 256 regular input + 768 cache_read + 512 cache_creation +
|
||||
// 200 output. Each input bucket bills at its own rate.
|
||||
cost, ok := tbl.Cost("anthropic", "claude-sonnet", 256, 200, 768, 512)
|
||||
require.True(t, ok, "known model resolves")
|
||||
want := (256.0/1000.0)*0.003 +
|
||||
(768.0/1000.0)*0.0003 +
|
||||
(512.0/1000.0)*0.00375 +
|
||||
(200.0/1000.0)*0.015
|
||||
assert.InDelta(t, want, cost, 1e-12,
|
||||
"each Anthropic input bucket must bill at its own configured rate")
|
||||
}
|
||||
|
||||
// TestCost_AnthropicCacheRatesFallBackToInput covers the no-opt-in
|
||||
// path: when neither CacheReadPer1K nor CacheCreationPer1K is set,
|
||||
// cache tokens bill at the regular input rate. This is more
|
||||
// accurate than today's behaviour (cache tokens ignored entirely)
|
||||
// without requiring operators to opt in via YAML.
|
||||
func TestCost_AnthropicCacheRatesFallBackToInput(t *testing.T) {
|
||||
tbl := &Table{entries: map[string]map[string]Entry{
|
||||
"anthropic": {"claude-sonnet": {InputPer1K: 0.003, OutputPer1K: 0.015}},
|
||||
}}
|
||||
cost, ok := tbl.Cost("anthropic", "claude-sonnet", 256, 200, 768, 512)
|
||||
require.True(t, ok)
|
||||
// Without overrides: every input bucket at input_per_1k.
|
||||
want := ((256.0+768.0+512.0)/1000.0)*0.003 + (200.0/1000.0)*0.015
|
||||
assert.InDelta(t, want, cost, 1e-12,
|
||||
"absent cache rates must fall back to input_per_1k — Anthropic cache tokens were ignored before this change, billing at input rate is more accurate as a default")
|
||||
}
|
||||
|
||||
func TestNewLoader_UnknownModel(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
|
||||
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, ok := l.Get().Cost("openai", "fantasy-model", 10, 10, 0, 0)
|
||||
assert.False(t, ok, "unknown model returns ok=false")
|
||||
|
||||
_, ok = l.Get().Cost("cohere", "anything", 10, 10, 0, 0)
|
||||
assert.False(t, ok, "unknown provider returns ok=false")
|
||||
}
|
||||
|
||||
func TestNewLoader_InvalidYAMLRejected(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
require.NoError(t, os.WriteFile(filepath.Join(base, "pricing.yaml"), []byte("\t- this is not: valid: yaml: :["), 0o600))
|
||||
|
||||
_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.Error(t, err, "invalid YAML must surface as construction error")
|
||||
}
|
||||
|
||||
func TestLoader_ReloadKeepsPreviousOnParseError(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
target := filepath.Join(base, "pricing.yaml")
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
|
||||
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, l.Get(), "initial table populated")
|
||||
|
||||
// Overwrite with content that violates the strict schema (extra field)
|
||||
// plus a bumped mtime to trigger reload.
|
||||
require.NoError(t, os.WriteFile(target, []byte("openai:\n gpt-4o:\n input_per_1k: 1.0\n output_per_1k: 2.0\n bogus_field: nope\n"), 0o600))
|
||||
future := time.Now().Add(time.Hour)
|
||||
require.NoError(t, os.Chtimes(target, future, future))
|
||||
|
||||
err = l.reload()
|
||||
require.Error(t, err, "parse error surfaced by reload()")
|
||||
|
||||
cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
|
||||
require.True(t, ok, "previous table still available after parse failure")
|
||||
assert.InDelta(t, 0.00075, cost, 1e-9, "previous cost preserved")
|
||||
}
|
||||
|
||||
func TestLoader_ReloadNoChangeIsNoOp(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
target := filepath.Join(base, "pricing.yaml")
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
|
||||
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err)
|
||||
ptrBefore := l.Get()
|
||||
|
||||
require.NoError(t, l.reload(), "no-change reload must not error")
|
||||
ptrAfter := l.Get()
|
||||
assert.Same(t, ptrBefore, ptrAfter, "table pointer unchanged when mtime unchanged")
|
||||
}
|
||||
|
||||
func TestLoader_ReloadDetectsChange(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
target := filepath.Join(base, "pricing.yaml")
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
|
||||
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
updated := []byte("openai:\n gpt-4o-mini:\n input_per_1k: 1.00\n output_per_1k: 2.00\n")
|
||||
require.NoError(t, os.WriteFile(target, updated, 0o600))
|
||||
future := time.Now().Add(time.Hour)
|
||||
require.NoError(t, os.Chtimes(target, future, future))
|
||||
|
||||
require.NoError(t, l.reload(), "reload must succeed on valid new content")
|
||||
|
||||
cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
|
||||
require.True(t, ok, "updated model still present")
|
||||
assert.InDelta(t, 3.0, cost, 0.0001, "new prices are applied: 1 + 2 per 1k")
|
||||
}
|
||||
|
||||
// TestLoader_ReloadGoroutinePicksUpChanges proves the background goroutine
|
||||
// started via Reload actually swaps the pricing table when the file changes
|
||||
// on disk. Without that goroutine running, pricing edits would never reach
|
||||
// requests until a proxy restart.
|
||||
func TestLoader_ReloadGoroutinePicksUpChanges(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
target := filepath.Join(base, "pricing.yaml")
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
|
||||
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err)
|
||||
l.SetReloadInterval(20 * time.Millisecond)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
l.Reload(ctx)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
// Before any rewrite, the loader holds the fixture's prices.
|
||||
costBefore, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
|
||||
require.True(t, ok, "fixture model must resolve initially")
|
||||
assert.InDelta(t, 0.00075, costBefore, 1e-9, "fixture prices apply before rewrite")
|
||||
|
||||
updated := []byte("openai:\n gpt-4o-mini:\n input_per_1k: 1.00\n output_per_1k: 2.00\n")
|
||||
require.NoError(t, os.WriteFile(target, updated, 0o600))
|
||||
future := time.Now().Add(time.Hour)
|
||||
require.NoError(t, os.Chtimes(target, future, future))
|
||||
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for {
|
||||
cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
|
||||
if ok && cost > 2.5 {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("background reloader did not pick up rewrite within deadline")
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
|
||||
cancel()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("Reload loop did not exit after cancel")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoader_ReloadBackgroundLoopCancellation(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
l.Reload(ctx)
|
||||
close(done)
|
||||
}()
|
||||
cancel()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("Reload loop did not exit on context cancel")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewLoader_PathValidation(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
relPath string
|
||||
}{
|
||||
{"traversal", "../../etc/passwd"},
|
||||
{"absolute", "/etc/passwd"},
|
||||
{"slash in basename", "sub/pricing.yaml"},
|
||||
{"control chars", "pricing\x00.yaml"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
_, err := NewLoader(base, tc.relPath, "llm_observability", nil)
|
||||
require.Error(t, err, "NewLoader must reject %q", tc.relPath)
|
||||
})
|
||||
}
|
||||
|
||||
// Empty relPath is no longer a validation error: the loader treats it
|
||||
// as "no override file, defaults only" so cost metadata is still
|
||||
// emitted for the embedded models out of the box.
|
||||
t.Run("empty falls back to defaults", func(t *testing.T) {
|
||||
l, err := NewLoader(base, "", "llm_observability", nil)
|
||||
require.NoError(t, err, "empty relPath should yield a defaults-only loader")
|
||||
require.NotNil(t, l, "loader must be returned")
|
||||
require.False(t, l.WatchesFile(), "no file watching when no override is given")
|
||||
_, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
|
||||
assert.True(t, ok, "embedded defaults should still resolve gpt-4o-mini")
|
||||
})
|
||||
}
|
||||
|
||||
// TestNewLoader_PathValidation_Extended covers the remaining attack shapes
|
||||
// called out in C2: dot references, embedded traversal segments, and a
|
||||
// newline in the basename. The basename regex must reject each one even
|
||||
// though filepath.Clean would otherwise collapse them.
|
||||
func TestNewLoader_PathValidation_Extended(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
relPath string
|
||||
}{
|
||||
{"dot", "."},
|
||||
{"dotdot", ".."},
|
||||
{"relative traversal", "../pricing.yaml"},
|
||||
{"embedded slash", "pri/cing.yaml"},
|
||||
{"newline", "pricing\n.yaml"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
_, err := NewLoader(base, tc.relPath, "llm_observability", nil)
|
||||
require.Error(t, err, "NewLoader must reject %q", tc.relPath)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewLoader_ValidBasenameLoads proves the allowlist is exclusive: a
|
||||
// basename containing only safe characters under baseDir loads. Without this
|
||||
// a regression that over-tightened the regex would silently break valid
|
||||
// deployments.
|
||||
func TestNewLoader_ValidBasenameLoads(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing-v2_prod.yaml"))
|
||||
|
||||
l, err := NewLoader(base, "pricing-v2_prod.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err, "basename with _, -, . must load")
|
||||
require.NotNil(t, l.Get(), "table populated")
|
||||
}
|
||||
|
||||
// TestNewLoader_SymlinkOutsideBaseDirRejected constructs a symlink under
|
||||
// baseDir that points to a file outside it. O_NOFOLLOW must refuse to open
|
||||
// the symlink even though the symlink path itself is a valid basename under
|
||||
// baseDir.
|
||||
func TestNewLoader_SymlinkOutsideBaseDirRejected(t *testing.T) {
|
||||
outside := t.TempDir()
|
||||
target := filepath.Join(outside, "evil.yaml")
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
|
||||
|
||||
base := t.TempDir()
|
||||
link := filepath.Join(base, "pricing.yaml")
|
||||
require.NoError(t, os.Symlink(target, link), "symlink setup")
|
||||
|
||||
_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.Error(t, err, "O_NOFOLLOW must reject symlink even when it points outside baseDir")
|
||||
}
|
||||
|
||||
func TestNewLoader_SymlinkRejected(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
concrete := filepath.Join(base, "real.yaml")
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), concrete)
|
||||
|
||||
link := filepath.Join(base, "pricing.yaml")
|
||||
require.NoError(t, os.Symlink(concrete, link), "symlink setup")
|
||||
|
||||
_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.Error(t, err, "O_NOFOLLOW must reject symlinked targets")
|
||||
}
|
||||
|
||||
func TestTableCost_NilSafe(t *testing.T) {
|
||||
var t1 *Table
|
||||
cost, ok := t1.Cost("x", "y", 1, 1, 0, 0)
|
||||
assert.False(t, ok, "nil table reports unknown")
|
||||
assert.Zero(t, cost, "nil table returns zero cost")
|
||||
assert.False(t, t1.Has("x", "y"), "nil table has nothing")
|
||||
}
|
||||
|
||||
func TestLoaderGet_NilSafe(t *testing.T) {
|
||||
var l *Loader
|
||||
assert.Nil(t, l.Get(), "nil loader returns nil table")
|
||||
}
|
||||
|
||||
// TestNewLoader_RejectsOversizedFile_FixesM4 proves the loader bounds reads
|
||||
// at maxPricingBytes so a hostile file cannot exhaust process memory.
|
||||
func TestNewLoader_RejectsOversizedFile_FixesM4(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
target := filepath.Join(base, "pricing.yaml")
|
||||
|
||||
// Build a YAML payload larger than the cap. We pad with valid YAML
|
||||
// comments so a partial read would still fail the size check rather
|
||||
// than the parser.
|
||||
header := "openai:\n"
|
||||
bigComment := make([]byte, maxPricingBytes+1024)
|
||||
for i := range bigComment {
|
||||
bigComment[i] = ' '
|
||||
}
|
||||
bigComment[0] = '#'
|
||||
bigComment[len(bigComment)-1] = '\n'
|
||||
payload := append([]byte(header), bigComment...)
|
||||
require.NoError(t, os.WriteFile(target, payload, 0o600))
|
||||
|
||||
_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.Error(t, err, "oversized pricing file must be rejected")
|
||||
assert.Contains(t, err.Error(), "exceeds", "rejection must reference the byte cap")
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
//go:build unix
|
||||
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// maxPricingBytes caps the size of the pricing YAML on read so a hostile or
|
||||
// runaway file cannot exhaust process memory during reload. 1 MiB is several
|
||||
// orders of magnitude larger than any reasonable pricing table.
|
||||
const maxPricingBytes int64 = 1 << 20
|
||||
|
||||
// loadPricing opens the file with O_NOFOLLOW, fstats the open descriptor,
|
||||
// and parses from that same descriptor. Never re-opens by path so a
|
||||
// mid-read rename or symlink swap cannot substitute content. Bytes are
|
||||
// capped at maxPricingBytes so the loader cannot be coerced into reading an
|
||||
// unbounded file.
|
||||
func loadPricing(path string) (*Table, time.Time, error) {
|
||||
f, err := os.OpenFile(path, os.O_RDONLY|syscall.O_NOFOLLOW, 0)
|
||||
if err != nil {
|
||||
return nil, time.Time{}, fmt.Errorf("open %s: %w", path, err)
|
||||
}
|
||||
defer func() {
|
||||
if cerr := f.Close(); cerr != nil {
|
||||
log.Debugf("close pricing file %s: %v", path, cerr)
|
||||
}
|
||||
}()
|
||||
|
||||
info, err := f.Stat()
|
||||
if err != nil {
|
||||
return nil, time.Time{}, fmt.Errorf("fstat %s: %w", path, err)
|
||||
}
|
||||
if !info.Mode().IsRegular() {
|
||||
return nil, time.Time{}, fmt.Errorf("pricing file %s is not a regular file", path)
|
||||
}
|
||||
|
||||
data, err := io.ReadAll(io.LimitReader(f, maxPricingBytes+1))
|
||||
if err != nil {
|
||||
return nil, time.Time{}, fmt.Errorf("read %s: %w", path, err)
|
||||
}
|
||||
if int64(len(data)) > maxPricingBytes {
|
||||
return nil, time.Time{}, fmt.Errorf("pricing file %s exceeds %d bytes", path, maxPricingBytes)
|
||||
}
|
||||
|
||||
table, err := parsePricingBytes(data)
|
||||
if err != nil {
|
||||
return nil, time.Time{}, err
|
||||
}
|
||||
return table, info.ModTime(), nil
|
||||
}
|
||||
|
||||
// statMtime returns the mtime of the file at path. It uses lstat semantics
|
||||
// via os.Lstat so a symlink swap is detected even though O_NOFOLLOW will
|
||||
// later reject the open.
|
||||
func statMtime(path string) (time.Time, error) {
|
||||
info, err := os.Lstat(path)
|
||||
if err != nil {
|
||||
return time.Time{}, fmt.Errorf("lstat %s: %w", path, err)
|
||||
}
|
||||
return info.ModTime(), nil
|
||||
}
|
||||
Reference in New Issue
Block a user