[management,proxy,client] 0.74.0 version (#6563)

* [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.

* Update getting started to point to rc when agent network enabled

* Add a reference to a commercial license

* Fix docs localhost link

* Fix docs localhost link

* Add private services domain note

* [management] Add agent-network telemetry metrics (#6561)

Surface agent-network adoption and usage in the self-hosted metrics
worker: distinct accounts, providers, policies, budget rules, accounts
with log collection enabled, and aggregated input/output tokens plus
cost.

Tokens and cost are summed from agent_network_request_usage (the
always-written per-request ledger) so the figures are accurate
regardless of the log-collection toggle and carry no double-counting.
All values come from a handful of indexed aggregate queries run only on
the worker's periodic tick.

Adds store.AgentNetworkMetrics with GetAgentNetworkMetrics on the Store
interface, the SqlStore implementation, and a zero-valued FileStore stub.

* Update NetBird server and proxy image versions to 0.74.0-rc.2

* [management,proxy] Reduce agent-network cognitive complexity (#6566)

Address the SonarCloud quality-gate findings in new agent-network code
by extracting focused helpers. No behavior change.

- synthesizer.go: split buildIdentityInjectConfigJSON into per-shape
  rule builders; extract mergeGuardrail from mergeGuardrails to cut
  nesting depth.
- llm_identity_inject: extract injectionEmitsAnything validation
  predicate from New.
- llm_response_parser/streaming.go: extract applyOpenAIStreamUsage and
  applyAnthropicStreamUsage (via a named anthropicStreamUsage type) and
  simplify the OpenAI scanner loop.
- reverseproxy.go: decompose ServeHTTP into serveRouteError,
  buildTargetContext, serveDirect, serveWithChain, captureRequestForChain,
  serveDeny, newResponseWriter, observeResponse, and forwardUpstream,
  preserving the defer ordering so response observation still reads the
  captured writer before it is released.

* [management] Move agent-network access-log ingest into the agentnetwork module (#6568)

The agent-network access-log ingest path (metaKey wire contract, flatten,
usage derivation, and the dual-write of the usage ledger + settings-gated
full row) lived in the reverseproxy accesslogs manager, even though the
agentnetwork module already owns the rest of that domain — types, read
(ListAccessLogs / GetUsageOverview), the budget-counter writes, and
retention cleanup.

Move it next to the rest: a stateless agentnetwork.IngestAccessLog(ctx,
store, entry) that the reverseproxy SaveAccessLog delegates to when the
entry is agent-network. Removes the agentNetworkTypes import from the
reverseproxy manager. No behavior change; the write/read table separation
is unchanged.

Adds real-store coverage for the disable->enable log-collection toggle
(usage ledger always written, full row gated) plus the metadata parse and
group-dedup helpers, which previously had no dedicated tests.

* Add session view support in the access log

* [management,proxy] Container-based agent-network e2e harness (#6577)

* [e2e] Add container-based agent-network e2e harness (Pillar 1)

Introduce a self-contained, OIDC-free e2e harness that stands up NetBird
in containers, so suites no longer depend on the hand-maintained Tilt
stack or a real IdP.

- harness brings up the combined server (management + signal + relay +
  STUN + embedded IdP) in a single container built from
  combined/Dockerfile.multistage, and mints an admin PAT through the
  unauthenticated /api/setup bootstrap (NB_SETUP_PAT_ENABLED). API access
  goes through the existing shared/management/client/rest typed client.
- the image is built via the docker CLI (BuildKit) so the Dockerfile's
  cache mounts are honored; testcontainers then runs the tagged image.
- everything is behind the `e2e` build tag so normal builds and unit
  tests never pull in testcontainers.

Adds BuildKit cache mounts to combined/Dockerfile.multistage so source
changes recompile incrementally rather than from scratch.

Pillar 1 proven by TestCombinedBootstrap: server builds, boots, mints a
PAT, and the PAT authenticates a real management API call.

* [e2e] Add management-side agent-network scenarios (Pillar 2)

Port the API-driven agent-network scenarios from the bash suites to Go,
sharing one combined server per package run (TestMain) with each test
owning its resource cleanup. Drives the /api/agent-network/* endpoints
through the shared REST client's NewRequest primitive with the generated
api types.

Scenarios:
- provider lifecycle (create/get/list/delete + 404 after delete)
- provider validation (missing api_key, unknown catalog id → 4xx)
- settings collection-toggle round-trip with cluster/subdomain immutability
- policy window floor (reject <60s enabled limit, accept at 60s)
- consumption read endpoint returns an array

All deterministic and dependency-free (dummy provider keys; no upstream
calls), so they run headless in CI.

* [e2e] Add live chat-through-proxy scenario (Pillar 3)

Stand up the full agent-network data path in containers and drive a real
chat-completion through the gateway:

- harness: a shared docker network (combined server reachable by alias),
  a proxy container built from the published reverse-proxy image
  (NB_PROXY_PRIVATE, NB_PROXY_ALLOW_INSECURE, NB_RELAY_TRANSPORT=ws to match
  the combined server's WS-multiplexed relay) with a generated self-signed
  wildcard cert, and a netbird client container that joins via a setup key.
- the combined image, proxy image, and client image default to the
  published rc.2 releases (overridable via NB_E2E_*_IMAGE; a bare local tag
  is built from source instead). Geolocation download is disabled so the
  server starts without external fetches.
- one shared domain is used for the management exposed address, the proxy
  domain, and the agent-network cluster; the proxy token is minted via the
  server CLI (global) to match the manual install.

TestChatCompletionThroughProxy provisions provider+policy+group+setup key,
runs proxy+client, drives an OpenAI chat-completion through the tunnel, and
asserts a 200 plus the ingested access-log row. Requires OPENAI_TOKEN
(skips otherwise). The provider must be created with enabled=true explicitly
— the create default is false despite the API doc.

* [e2e] Run the live chat scenario across a provider matrix

Replace the single-provider chat test with a data-driven matrix that runs
the same scenario through every provider whose credentials are present in
the environment (keys/URLs sourced from ~/.llm-keys locally, Actions
secrets in CI):

- OpenAI (chat), Anthropic (messages), Vercel, OpenRouter, Cloudflare
  (OpenAI-compatible gateways), and Bedrock (path-routed, bearer, via the
  messages shape) — covering both wire shapes and the gateway routing.
- all providers are created enabled with a unique model string so the
  proxy's connect-time snapshot carries them all and model->provider
  routing is unambiguous (provider toggles after connect don't reconcile
  to a connected proxy).
- the client supports both wire shapes (/v1/chat/completions and
  /v1/messages); Cloudflare gets the openai provider segment appended to
  its gateway URL.

Each provider must return 200 through the tunnel and produce an ingested
access-log row. Vertex is intentionally excluded from the uniform matrix:
it needs a bespoke rawPredict request shape rather than the shared
chat/messages path, so it warrants a dedicated scenario.

* [ci] Add manual workflow for the agent-network e2e suite

The e2e suite (build tag `e2e`) stands up the combined server + proxy +
client in Docker and drives live chat-completions, so it is slow and needs
provider credentials. Gate it out of normal CI (it already is, via the
build tag) and run it on demand via workflow_dispatch. Provider scenarios
skip when their secret is unset, so it degrades gracefully.

* [e2e] Add Vertex to the provider matrix; run e2e on ubuntu-latest

Vertex (Anthropic-on-Vertex) doesn't share the chat/messages wire shapes:
the model travels in a rawPredict path and the proxy mints the service
account's OAuth token. Add a Vertex client method that posts
/v1/projects/<project>/locations/<region>/publishers/anthropic/models/<model>:rawPredict
with the Vertex anthropic_version body, and wire it into the matrix as a
path-routed provider (created without a models array). It is keyed off
GOOGLE_VERTEX_SA_BASE64 + GOOGLE_VERTEX_PROJECT (region defaults to
"global", model to a pinned claude snapshot, both overridable).

Also bump the e2e workflow runner to ubuntu-latest and add the Vertex
secrets.

* Add docker/docker and docker/go-connections as direct dependencies in go.mod

* [ci] Trigger agent-network e2e workflow on push to main and pull requests

* [e2e] Fix proxy cert permission denied on Linux CI runners

The proxy bind-mounts a temp dir of self-signed certs. MkdirTemp creates
it 0700 and the key was 0600, which Docker Desktop on macOS ignores but a
non-root proxy container on Linux runners cannot traverse/read, so the
cert watcher failed with "open /certs/tls.crt: permission denied" and the
container exited. Widen the cert dir to 0755 and write the throwaway key
0644 so the proxy uid can read the bind-mounted material.

* [e2e] Build images from source by default instead of pulling rc.2

The agent-network code under test lives in this branch, so the e2e should
exercise it rather than a frozen published release. Flip the harness
default: combined/proxy/client are now built from their in-repo
Dockerfiles (combined/Dockerfile.multistage, proxy/Dockerfile.multistage,
e2e/harness/Dockerfile.client) under local tags. Pulling a published image
stays available by setting NB_E2E_*_IMAGE to a registry reference.

Builds now go through buildx --load so the Dockerfile cache mounts are
honored and the result is loaded for testcontainers. The CI workflow adds
a container-driver builder and a local layer cache (NB_E2E_BUILDX_CACHE)
persisted via actions/cache, which caches the base/apt/dep-download layers
across runs. The Go compile still re-runs each time, as BuildKit mount
caches cannot be exported to the GitHub cache.

* [e2e] Cover real providers in lifecycle + assert real consumption metering

- TestProviderLifecycle now runs per available real provider (create → get →
  list → delete → 404) instead of a single dummy provider, exercising each
  catalog's create and field round-trip. Create is offline, so it stays fast
  and burns no provider quota; falls back to a synthetic OpenAI provider when
  no keys are set.
- TestProvidersMatrix attaches a token limit (high caps, 60s window) to its
  policy, which switches on usage metering, and asserts consumption rows are
  recorded with positive token counts after the live traffic. Consumption is
  account-scoped (keyed by source group / user and window, not per provider),
  so the assertion is aggregate.
- TestProviderValidation gains invalid-upstream and blank-name cases. Create
  validation is uniform across catalogs (no per-provider required-field rules),
  so per-provider rejection cases would be redundant.

* [e2e] Assert session id propagates per provider

Each matrix request now sends a unique session id as the universal
x-session-id header and asserts it round-trips into that provider's
access-log row. This guards the session-grouping contract end to end for
every provider (header extraction runs in llm_request_parser ahead of the
parser-specific body extraction, so it is provider-agnostic).

* [e2e] Drop accidentally committed sync-phases dashboard

netbird-sync-phases.json was swept into the Pillar 1 commit by a broad
git add; it belongs to the unrelated sync-phases metrics work, not this
e2e harness. Remove it from the branch so the PR diff is scoped to the
e2e changes.

* [e2e] Revert accidentally committed sync-phase ingest spec

The netbird_sync_phase measurement spec in metrics ingest was swept into
the Pillar 1 commit; it belongs to the unrelated sync-phases metrics work,
not this e2e harness. Its emission side never landed here, so the spec was
orphaned anyway. Restore ingest/main.go to its origin/main state.

* Fix golint issues

* Fix sonar

* Add access log session test

* Fix access log tests

---------

Co-authored-by: braginini <bangvalo@gmail.com>
Co-authored-by: Zoltan Papp <zoltan.pmail@gmail.com>
This commit is contained in:
Maycon Santos
2026-07-01 12:45:14 +02:00
committed by GitHub
co-authored by braginini Zoltan Papp
parent 3be90f06b2
commit 92a66cdd20
208 changed files with 39957 additions and 688 deletions
+196
View File
@@ -0,0 +1,196 @@
package llm
import (
"encoding/json"
"fmt"
"strings"
)
// AnthropicParser implements the Parser interface for the Anthropic Messages
// and Completions APIs. Detection is substring-based to tolerate upstream
// path rewrites.
type AnthropicParser struct{}
var anthropicPathHints = []string{
"/v1/messages",
"/v1/complete",
}
// Provider returns ProviderAnthropic.
func (AnthropicParser) Provider() Provider { return ProviderAnthropic }
// ProviderName returns the stable label used for metrics and metadata.
func (AnthropicParser) ProviderName() string { return "anthropic" }
// DetectFromURL reports whether the given request path looks like an
// Anthropic API endpoint. The match is case-insensitive and substring-based.
func (AnthropicParser) DetectFromURL(path string) bool {
lower := strings.ToLower(path)
for _, hint := range anthropicPathHints {
if strings.Contains(lower, hint) {
return true
}
}
return false
}
type anthropicRequest struct {
Model string `json:"model"`
Stream *bool `json:"stream"`
System json.RawMessage `json:"system"`
Messages []anthropicMessage `json:"messages"`
// Legacy /v1/complete endpoint.
Prompt string `json:"prompt"`
}
type anthropicMessage struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
}
// ParseRequest extracts the model name and streaming flag from an Anthropic
// request body. Unknown or missing fields leave the corresponding struct
// members zero-valued.
func (AnthropicParser) ParseRequest(body []byte) (RequestFacts, error) {
var req anthropicRequest
if err := json.Unmarshal(body, &req); err != nil {
return RequestFacts{}, fmt.Errorf("decode anthropic request: %w: %v", ErrMalformedRequest, err)
}
return RequestFacts{
Model: req.Model,
Stream: ptrDeref(req.Stream),
}, nil
}
type anthropicResponse struct {
Usage struct {
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
// CacheReadInputTokens and CacheCreationInputTokens are
// ADDITIVE to InputTokens (not subset), each billed at its
// own rate by the cost meter. cache_read is the cheaper
// read-from-cache rate, cache_creation is the more
// expensive write-to-cache rate.
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
} `json:"usage"`
}
// ParseResponse decodes the non-streaming Anthropic response envelope. Status
// codes other than 200 are treated as non-LLM responses so the caller can
// skip cost accounting without aborting the request.
func (AnthropicParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) {
if status != 200 {
return Usage{}, fmt.Errorf("anthropic status %d: %w", status, ErrNotLLMResponse)
}
if isEventStream(contentType) {
return Usage{}, ErrStreamingUnsupported
}
if !isJSON(contentType) {
return Usage{}, fmt.Errorf("anthropic content-type %q: %w", contentType, ErrNotLLMResponse)
}
var resp anthropicResponse
if err := json.Unmarshal(body, &resp); err != nil {
return Usage{}, fmt.Errorf("decode anthropic response: %w: %v", ErrMalformedResponse, err)
}
return Usage{
InputTokens: resp.Usage.InputTokens,
OutputTokens: resp.Usage.OutputTokens,
TotalTokens: resp.Usage.InputTokens + resp.Usage.OutputTokens + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens,
CachedInputTokens: resp.Usage.CacheReadInputTokens,
CacheCreationTokens: resp.Usage.CacheCreationInputTokens,
}, nil
}
// ExtractPrompt returns the user-visible prompt text from an Anthropic
// request body. Handles the Messages API (system + messages[]) and the
// legacy /v1/complete prompt string. Returns "" on any decode failure.
func (AnthropicParser) ExtractPrompt(body []byte) string {
var req anthropicRequest
if err := json.Unmarshal(body, &req); err != nil {
return ""
}
var b strings.Builder
if len(req.System) > 0 {
if s := decodeStringOrJoin(req.System); s != "" {
b.WriteString("system: ")
b.WriteString(s)
}
}
for _, m := range req.Messages {
if b.Len() > 0 {
b.WriteByte('\n')
}
if m.Role != "" {
b.WriteString(m.Role)
b.WriteString(": ")
}
b.WriteString(decodeStringOrJoin(m.Content))
}
if b.Len() == 0 && req.Prompt != "" {
b.WriteString(req.Prompt)
}
return b.String()
}
// ExtractSessionID is the body-side fallback for Anthropic. Claude Code's
// authoritative session marker is the X-Claude-Code-Session-Id request
// header (handled by the request-parser middleware); this only mines the
// optional metadata.user_id for an embedded "...session_<uuid>" marker.
// metadata.user_id on its own is a USER identifier, not a session, so the
// whole value is deliberately NOT used — returning it would mislabel every
// request from a user as one session. Returns "" when no session marker is
// present.
func (AnthropicParser) ExtractSessionID(body []byte) string {
var req struct {
Metadata struct {
UserID string `json:"user_id"`
} `json:"metadata"`
}
if err := json.Unmarshal(body, &req); err != nil {
return ""
}
if idx := strings.LastIndex(req.Metadata.UserID, "session_"); idx >= 0 {
if session := req.Metadata.UserID[idx+len("session_"):]; session != "" {
return session
}
}
return ""
}
type anthropicMessageResponse struct {
Content []struct {
Type string `json:"type"`
Text string `json:"text"`
} `json:"content"`
// Legacy /v1/complete response.
Completion string `json:"completion"`
}
// ExtractCompletion returns the assistant text from a non-streaming Anthropic
// Messages or Completions response. Returns "" when status/content-type
// indicate the body is not parseable or no text part is present.
func (AnthropicParser) ExtractCompletion(status int, contentType string, body []byte) string {
if status != 200 || isEventStream(contentType) || !isJSON(contentType) {
return ""
}
var resp anthropicMessageResponse
if err := json.Unmarshal(body, &resp); err != nil {
return ""
}
var b strings.Builder
for _, part := range resp.Content {
if part.Text == "" {
continue
}
if b.Len() > 0 {
b.WriteByte('\n')
}
b.WriteString(part.Text)
}
if b.Len() == 0 {
return resp.Completion
}
return b.String()
}
+169
View File
@@ -0,0 +1,169 @@
package llm
import (
"errors"
"os"
"path/filepath"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAnthropicDetectFromURL(t *testing.T) {
p := AnthropicParser{}
cases := map[string]bool{
"/v1/messages": true,
"/v1/complete": true,
"/V1/Messages": true,
"/proxy/v1/messages?x": true,
"/v1/chat/completions": false,
"": false,
}
for path, want := range cases {
assert.Equal(t, want, p.DetectFromURL(path), "DetectFromURL(%q)", path)
}
}
func TestAnthropicParseRequest(t *testing.T) {
p := AnthropicParser{}
t.Run("stream true", func(t *testing.T) {
facts, err := p.ParseRequest([]byte(`{"model":"claude-sonnet-4-5","stream":true}`))
require.NoError(t, err)
assert.Equal(t, "claude-sonnet-4-5", facts.Model, "model extracted")
assert.True(t, facts.Stream, "stream flag honoured")
})
t.Run("stream default", func(t *testing.T) {
facts, err := p.ParseRequest([]byte(`{"model":"claude-sonnet-4-5"}`))
require.NoError(t, err)
assert.False(t, facts.Stream, "missing stream flag defaults to false")
})
t.Run("malformed", func(t *testing.T) {
_, err := p.ParseRequest([]byte(`{"model":`))
require.Error(t, err)
assert.True(t, errors.Is(err, ErrMalformedRequest), "sentinel wrapped")
})
}
func TestAnthropicParseResponse(t *testing.T) {
p := AnthropicParser{}
t.Run("happy fixture", func(t *testing.T) {
body, err := os.ReadFile(filepath.Join("fixtures", "anthropic_messages.json"))
require.NoError(t, err, "fixture must be readable")
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(123), usage.InputTokens, "input tokens extracted")
assert.Equal(t, int64(45), usage.OutputTokens, "output tokens extracted")
assert.Equal(t, int64(168), usage.TotalTokens, "total computed as sum")
})
t.Run("streaming rejected", func(t *testing.T) {
_, err := p.ParseResponse(200, "text/event-stream", []byte(""))
require.ErrorIs(t, err, ErrStreamingUnsupported, "SSE responses must use the scanner")
})
t.Run("non-200", func(t *testing.T) {
_, err := p.ParseResponse(429, "application/json", []byte(`{}`))
require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 rejected as non-LLM")
})
t.Run("non-json content type", func(t *testing.T) {
_, err := p.ParseResponse(200, "text/html", []byte(`{}`))
require.ErrorIs(t, err, ErrNotLLMResponse, "text/html treated as non-LLM")
})
t.Run("malformed body", func(t *testing.T) {
_, err := p.ParseResponse(200, "application/json", []byte(`{`))
require.ErrorIs(t, err, ErrMalformedResponse, "bad JSON yields malformed error")
})
// Anthropic's two cache fields are ADDITIVE to input_tokens (not
// subset). The parser must surface them so the cost meter can
// bill each bucket at its own configured rate. Total includes
// every bucket so downstream attribution sees the full token
// volume the request consumed.
t.Run("cache_read_input_tokens surfaces as CachedInputTokens (additive)", func(t *testing.T) {
body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(256), usage.InputTokens, "regular input remains separate from cache buckets")
assert.Equal(t, int64(768), usage.CachedInputTokens, "cache_read maps onto CachedInputTokens — same field carries OpenAI cached subset and Anthropic cache reads")
assert.Zero(t, usage.CacheCreationTokens)
assert.Equal(t, int64(256+200+768), usage.TotalTokens, "total includes every input bucket plus output — cache reads are billable tokens")
})
t.Run("cache_creation_input_tokens surfaces as CacheCreationTokens (additive)", func(t *testing.T) {
body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_creation_input_tokens":512}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(256), usage.InputTokens)
assert.Zero(t, usage.CachedInputTokens)
assert.Equal(t, int64(512), usage.CacheCreationTokens, "cache_creation surfaces — meter applies the write-rate multiplier")
assert.Equal(t, int64(256+200+512), usage.TotalTokens)
})
t.Run("both cache buckets present", func(t *testing.T) {
body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768,"cache_creation_input_tokens":512}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(768), usage.CachedInputTokens)
assert.Equal(t, int64(512), usage.CacheCreationTokens)
assert.Equal(t, int64(256+200+768+512), usage.TotalTokens, "all four buckets sum into total")
})
t.Run("absent cache fields leave counts at zero", func(t *testing.T) {
body := []byte(`{"usage":{"input_tokens":100,"output_tokens":50}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Zero(t, usage.CachedInputTokens, "no cache_read field = no cached count")
assert.Zero(t, usage.CacheCreationTokens, "no cache_creation field = no creation count")
assert.Equal(t, int64(150), usage.TotalTokens, "back to the simple in+out total when no cache buckets present")
})
}
func TestAnthropicExtractPrompt_Messages(t *testing.T) {
body := []byte(`{"model":"claude-sonnet-4-7","system":"be brief","messages":[{"role":"user","content":"hi"},{"role":"assistant","content":"yes"}]}`)
got := AnthropicParser{}.ExtractPrompt(body)
require.Contains(t, got, "system: be brief", "system surfaces with role label")
require.Contains(t, got, "user: hi", "user message surfaces")
require.Contains(t, got, "assistant: yes", "assistant message surfaces")
}
func TestAnthropicExtractPrompt_LegacyComplete(t *testing.T) {
body := []byte(`{"model":"claude-2","prompt":"\n\nHuman: hi\n\nAssistant:"}`)
got := AnthropicParser{}.ExtractPrompt(body)
require.Contains(t, got, "Human: hi", "legacy prompt string surfaces")
}
func TestAnthropicExtractSessionID(t *testing.T) {
t.Run("claude code session suffix", func(t *testing.T) {
body := []byte(`{"model":"claude-opus-4-8","metadata":{"user_id":"user_abc123_account_def456_session_9f8e7d6c"},"messages":[]}`)
assert.Equal(t, "9f8e7d6c", AnthropicParser{}.ExtractSessionID(body), "session_<id> suffix must be extracted from metadata.user_id")
})
t.Run("plain user_id is not treated as a session", func(t *testing.T) {
body := []byte(`{"model":"claude-opus-4-8","metadata":{"user_id":"acme-team"},"messages":[]}`)
assert.Equal(t, "", AnthropicParser{}.ExtractSessionID(body), "a user identifier without a session marker must NOT be used as a session id")
})
t.Run("no metadata yields empty", func(t *testing.T) {
body := []byte(`{"model":"claude-opus-4-8","messages":[{"role":"user","content":"hi"}]}`)
assert.Equal(t, "", AnthropicParser{}.ExtractSessionID(body), "absent metadata.user_id yields no session id")
})
}
func TestAnthropicExtractCompletion_Messages(t *testing.T) {
body, err := os.ReadFile(filepath.Join("fixtures", "anthropic_messages.json"))
require.NoError(t, err)
got := AnthropicParser{}.ExtractCompletion(200, "application/json", body)
require.NotEmpty(t, got, "anthropic fixture has assistant text")
}
func TestAnthropicExtractCompletion_Streaming(t *testing.T) {
got := AnthropicParser{}.ExtractCompletion(200, "text/event-stream", []byte(""))
require.Empty(t, got, "streaming responses are skipped")
}
+189
View File
@@ -0,0 +1,189 @@
package llm
import (
"encoding/json"
"fmt"
"strings"
)
// ProviderNameBedrock is the stable label for the AWS Bedrock parser, used as
// the llm.provider metadata value and the cost-meter formula selector.
const ProviderNameBedrock = "bedrock"
// BedrockParser implements the Parser interface for the AWS Bedrock runtime.
// Bedrock carries the model in the URL path (/model/{id}/{action}); the request
// middleware extracts it there, so this parser focuses on the response shapes:
// the vendor-native InvokeModel body (e.g. Anthropic's snake_case usage) and the
// unified Converse body (camelCase usage).
type BedrockParser struct{}
var bedrockPathHints = []string{"/invoke", "/converse"}
// Provider returns ProviderBedrock.
func (BedrockParser) Provider() Provider { return ProviderBedrock }
// ProviderName returns the stable label used for metrics and metadata.
func (BedrockParser) ProviderName() string { return ProviderNameBedrock }
// DetectFromURL reports whether the path is a Bedrock runtime model endpoint.
func (BedrockParser) DetectFromURL(path string) bool {
lower := strings.ToLower(path)
if !strings.HasPrefix(lower, "/model/") {
return false
}
for _, hint := range bedrockPathHints {
if strings.Contains(lower, hint) {
return true
}
}
return false
}
// ParseRequest is a no-op for Bedrock: the model lives in the URL path, not the
// body, and the streaming flag is derived from the path action. The request
// middleware handles both via parseBedrockPath, so this returns empty facts.
func (BedrockParser) ParseRequest([]byte) (RequestFacts, error) {
return RequestFacts{}, nil
}
// bedrockResponse captures token usage from both Bedrock response shapes:
// InvokeModel (vendor-native; Anthropic uses snake_case + additive cache
// buckets) and Converse (camelCase, with a precomputed total).
type bedrockResponse struct {
Usage struct {
// InvokeModel (Anthropic-on-Bedrock) — snake_case.
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
// Converse — camelCase.
InputTokensCamel int64 `json:"inputTokens"`
OutputTokensCamel int64 `json:"outputTokens"`
TotalTokensCamel int64 `json:"totalTokens"`
} `json:"usage"`
}
// ParseResponse decodes the non-streaming Bedrock response envelope, handling
// both the InvokeModel and Converse usage shapes. Non-200 / non-JSON bodies are
// treated as non-LLM responses so the caller skips cost accounting.
func (BedrockParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) {
if status != 200 {
return Usage{}, fmt.Errorf("bedrock status %d: %w", status, ErrNotLLMResponse)
}
if isAWSEventStream(contentType) || isEventStream(contentType) {
return Usage{}, ErrStreamingUnsupported
}
if !isJSON(contentType) {
return Usage{}, fmt.Errorf("bedrock content-type %q: %w", contentType, ErrNotLLMResponse)
}
var resp bedrockResponse
if err := json.Unmarshal(body, &resp); err != nil {
return Usage{}, fmt.Errorf("decode bedrock response: %w: %v", ErrMalformedResponse, err)
}
inTok := firstNonZero(resp.Usage.InputTokens, resp.Usage.InputTokensCamel)
outTok := firstNonZero(resp.Usage.OutputTokens, resp.Usage.OutputTokensCamel)
total := resp.Usage.TotalTokensCamel
if total == 0 {
total = inTok + outTok + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens
}
return Usage{
InputTokens: inTok,
OutputTokens: outTok,
TotalTokens: total,
CachedInputTokens: resp.Usage.CacheReadInputTokens,
CacheCreationTokens: resp.Usage.CacheCreationInputTokens,
}, nil
}
// ExtractPrompt returns the user-visible prompt from a Bedrock request body,
// handling both the InvokeModel (Anthropic Messages: system + messages[]) and
// Converse (messages[].content[].text) shapes. Returns "" on decode failure.
func (BedrockParser) ExtractPrompt(body []byte) string {
var req struct {
System json.RawMessage `json:"system"`
Messages []struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
} `json:"messages"`
}
if err := json.Unmarshal(body, &req); err != nil {
return ""
}
var b strings.Builder
if s := decodeStringOrJoin(req.System); s != "" {
b.WriteString("system: ")
b.WriteString(s)
}
for _, m := range req.Messages {
if b.Len() > 0 {
b.WriteByte('\n')
}
if m.Role != "" {
b.WriteString(m.Role)
b.WriteString(": ")
}
b.WriteString(decodeStringOrJoin(m.Content))
}
return b.String()
}
// ExtractCompletion returns the assistant text from a non-streaming Bedrock
// response, handling InvokeModel (Anthropic content[].text) and Converse
// (output.message.content[].text).
func (BedrockParser) ExtractCompletion(status int, contentType string, body []byte) string {
if status != 200 || isAWSEventStream(contentType) || !isJSON(contentType) {
return ""
}
var resp struct {
Content []struct {
Text string `json:"text"`
} `json:"content"`
Output struct {
Message struct {
Content []struct {
Text string `json:"text"`
} `json:"content"`
} `json:"message"`
} `json:"output"`
}
if err := json.Unmarshal(body, &resp); err != nil {
return ""
}
var b strings.Builder
appendText := func(text string) {
if text == "" {
return
}
if b.Len() > 0 {
b.WriteByte('\n')
}
b.WriteString(text)
}
for _, p := range resp.Content {
appendText(p.Text)
}
for _, p := range resp.Output.Message.Content {
appendText(p.Text)
}
return b.String()
}
// ExtractSessionID has no Bedrock-native marker; session grouping relies on the
// request headers handled by the middleware. Returns "".
func (BedrockParser) ExtractSessionID([]byte) string { return "" }
// firstNonZero returns a when non-zero, else b. Folds the snake_case and
// camelCase usage variants into a single value.
func firstNonZero(a, b int64) int64 {
if a != 0 {
return a
}
return b
}
// isAWSEventStream reports whether contentType is the AWS binary event-stream
// framing used by Bedrock's streaming endpoints.
func isAWSEventStream(contentType string) bool {
return strings.Contains(strings.ToLower(contentType), "application/vnd.amazon.eventstream")
}
+65
View File
@@ -0,0 +1,65 @@
package llm
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestBedrockParser_ParseResponse_Invoke(t *testing.T) {
body := []byte(`{"usage":{"input_tokens":13,"output_tokens":5,"cache_read_input_tokens":2,"cache_creation_input_tokens":4}}`)
u, err := BedrockParser{}.ParseResponse(200, "application/json", body)
require.NoError(t, err)
require.Equal(t, int64(13), u.InputTokens, "invoke input tokens")
require.Equal(t, int64(5), u.OutputTokens, "invoke output tokens")
require.Equal(t, int64(2), u.CachedInputTokens, "invoke cache-read tokens")
require.Equal(t, int64(4), u.CacheCreationTokens, "invoke cache-creation tokens")
require.Equal(t, int64(13+5+2+4), u.TotalTokens, "invoke total is additive")
}
func TestBedrockParser_ParseResponse_Converse(t *testing.T) {
body := []byte(`{"output":{"message":{"content":[{"text":"pong"}]}},"usage":{"inputTokens":11,"outputTokens":3,"totalTokens":14}}`)
u, err := BedrockParser{}.ParseResponse(200, "application/json", body)
require.NoError(t, err)
require.Equal(t, int64(11), u.InputTokens, "converse camelCase input tokens")
require.Equal(t, int64(3), u.OutputTokens, "converse camelCase output tokens")
require.Equal(t, int64(14), u.TotalTokens, "converse uses provider total")
}
func TestBedrockParser_ParseResponse_StreamingUnsupported(t *testing.T) {
_, err := BedrockParser{}.ParseResponse(200, "application/vnd.amazon.eventstream", []byte("binary"))
require.ErrorIs(t, err, ErrStreamingUnsupported, "event-stream must route to the streaming accumulator")
}
func TestBedrockParser_ParseResponse_NonSuccess(t *testing.T) {
_, err := BedrockParser{}.ParseResponse(404, "application/json", []byte(`{"message":"gated"}`))
require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 is not an LLM response")
}
func TestBedrockParser_ExtractCompletion(t *testing.T) {
invoke := BedrockParser{}.ExtractCompletion(200, "application/json", []byte(`{"content":[{"text":"a"},{"text":"b"}]}`))
require.Equal(t, "a\nb", invoke, "invoke completion joins content parts")
converse := BedrockParser{}.ExtractCompletion(200, "application/json", []byte(`{"output":{"message":{"content":[{"text":"x"}]}}}`))
require.Equal(t, "x", converse, "converse completion reads output.message.content")
}
func TestBedrockParser_ExtractPrompt(t *testing.T) {
invoke := BedrockParser{}.ExtractPrompt([]byte(`{"messages":[{"role":"user","content":"hi"}]}`))
require.Equal(t, "user: hi", invoke, "invoke prompt reads anthropic content string")
converse := BedrockParser{}.ExtractPrompt([]byte(`{"messages":[{"role":"user","content":[{"text":"hello"}]}]}`))
require.Equal(t, "user: hello", converse, "converse prompt reads content parts")
}
func TestBedrockParser_DetectFromURL(t *testing.T) {
require.True(t, BedrockParser{}.DetectFromURL("/model/eu.anthropic.claude/invoke"), "invoke path")
require.True(t, BedrockParser{}.DetectFromURL("/model/x/converse-stream"), "converse-stream path")
require.False(t, BedrockParser{}.DetectFromURL("/v1/chat/completions"), "openai path is not bedrock")
}
func TestBedrockParser_RegisteredByName(t *testing.T) {
p, ok := ParserByName(ProviderNameBedrock)
require.True(t, ok, "bedrock parser is registered")
require.Equal(t, ProviderNameBedrock, p.ProviderName())
}
+31
View File
@@ -0,0 +1,31 @@
package llm
import "errors"
// Sentinel errors returned by parsers and the pricing loader. Callers use
// errors.Is to branch on a condition without coupling to parser internals.
var (
// ErrUnknownProvider indicates no parser claimed the request path.
ErrUnknownProvider = errors.New("llmobs: unknown provider")
// ErrUnsupportedModel indicates the response parsed successfully but the
// model is absent from the pricing table. Token counts are still valid.
ErrUnsupportedModel = errors.New("llmobs: unsupported model")
// ErrNotLLMResponse indicates the response is not a JSON success body
// that a non-streaming parser can consume (non-200 or wrong content type).
ErrNotLLMResponse = errors.New("llmobs: not an LLM response")
// ErrStreamingUnsupported indicates the caller passed an SSE response to
// a non-streaming parser. Streaming is handled separately via the SSE
// scanner.
ErrStreamingUnsupported = errors.New("llmobs: streaming response requires SSE scanner")
// ErrMalformedResponse indicates the response body could not be decoded
// as the provider-specific JSON schema.
ErrMalformedResponse = errors.New("llmobs: malformed response body")
// ErrMalformedRequest indicates the request body could not be decoded as
// the provider-specific JSON schema.
ErrMalformedRequest = errors.New("llmobs: malformed request body")
)
@@ -0,0 +1,17 @@
{
"id": "msg_abc",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5",
"content": [
{
"type": "text",
"text": "Hello, world!"
}
],
"stop_reason": "end_turn",
"usage": {
"input_tokens": 123,
"output_tokens": 45
}
}
@@ -0,0 +1,21 @@
event: message_start
data: {"type":"message_start","message":{"id":"msg_abc","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[],"stop_reason":null,"usage":{"input_tokens":123,"output_tokens":1}}}
event: content_block_start
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":", world!"}}
event: content_block_stop
data: {"type":"content_block_stop","index":0}
event: message_delta
data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":45}}
event: message_stop
data: {"type":"message_stop"}
@@ -0,0 +1,21 @@
{
"id": "chatcmpl-abc",
"object": "chat.completion",
"created": 1700000000,
"model": "gpt-4o-mini",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello, world!"
},
"finish_reason": "stop"
}
],
"usage": {
"prompt_tokens": 123,
"completion_tokens": 45,
"total_tokens": 168
}
}
@@ -0,0 +1,24 @@
{
"id": "resp_abc",
"object": "response",
"created_at": 1700000000,
"model": "gpt-5.4",
"output": [
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": "ok"}]
}
],
"usage": {
"input_tokens": 15,
"input_tokens_details": {
"cached_tokens": 0
},
"output_tokens": 414,
"output_tokens_details": {
"reasoning_tokens": 0
},
"total_tokens": 429
}
}
@@ -0,0 +1,24 @@
event: response.created
data: {"type":"response.created","response":{"id":"resp_abc","object":"response","model":"gpt-5.5","usage":null}}
event: response.in_progress
data: {"type":"response.in_progress","response":{"id":"resp_abc","usage":null}}
event: response.output_item.added
data: {"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant","content":[]}}
event: response.content_part.added
data: {"type":"response.content_part.added","item_id":"msg_1","output_index":0,"content_index":0,"part":{"type":"output_text","text":""}}
event: response.output_text.delta
data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"Hello"}
event: response.output_text.delta
data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":", world!"}
event: response.output_text.done
data: {"type":"response.output_text.done","item_id":"msg_1","output_index":0,"content_index":0,"text":"Hello, world!"}
event: response.completed
data: {"type":"response.completed","response":{"id":"resp_abc","object":"response","model":"gpt-5.5","usage":{"input_tokens":123,"input_tokens_details":{"cached_tokens":40},"output_tokens":45,"output_tokens_details":{"reasoning_tokens":12},"total_tokens":168}}}
@@ -0,0 +1,8 @@
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":", world!"},"finish_reason":null}]}
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":123,"completion_tokens":45,"total_tokens":168}}
data: [DONE]
+59
View File
@@ -0,0 +1,59 @@
# Realistic-pricing starter for llm_observability. Drop this into the
# directory you point the proxy at via --plugin-data-dir, then reference it
# from the target's plugin config:
#
# plugins:
# - id: llm_observability
# enabled: true
# params:
# pricing_path: pricing.yaml
#
# Values are USD per 1_000 tokens. Public list prices drift; treat this as a
# starting point and keep your production copy current.
openai:
# GPT-5 family
gpt-5:
input_per_1k: 0.00125
output_per_1k: 0.01
gpt-5-mini:
input_per_1k: 0.00025
output_per_1k: 0.002
gpt-5-nano:
input_per_1k: 0.00005
output_per_1k: 0.0004
gpt-5.4:
input_per_1k: 0.00125
output_per_1k: 0.01
# GPT-4o family
gpt-4o:
input_per_1k: 0.0025
output_per_1k: 0.01
gpt-4o-mini:
input_per_1k: 0.00015
output_per_1k: 0.0006
# Embeddings
text-embedding-3-large:
input_per_1k: 0.00013
output_per_1k: 0
text-embedding-3-small:
input_per_1k: 0.00002
output_per_1k: 0
anthropic:
# Claude 4.x family
claude-opus-4-7:
input_per_1k: 0.015
output_per_1k: 0.075
claude-sonnet-4-7:
input_per_1k: 0.003
output_per_1k: 0.015
claude-sonnet-4-6:
input_per_1k: 0.003
output_per_1k: 0.015
claude-sonnet-4-5:
input_per_1k: 0.003
output_per_1k: 0.015
claude-haiku-4-5:
input_per_1k: 0.0008
output_per_1k: 0.004
+412
View File
@@ -0,0 +1,412 @@
package llm
import (
"encoding/json"
"fmt"
"strings"
)
// OpenAIParser implements the Parser interface for OpenAI-compatible APIs.
// It recognizes chat.completions, completions, embeddings, and the newer
// responses endpoint; any proxy path-prefix stripping is tolerated by the
// substring match in DetectFromURL.
type OpenAIParser struct{}
// openAIPathHints are substring patterns that mark a request as
// OpenAI-shaped. The bare `/chat/completions` is listed alongside
// `/v1/chat/completions` because gateways like Cloudflare AI
// Gateway place their own version segment before the provider
// slug (gateway/v1/{account}/{gateway}/openai/chat/completions) —
// the canonical `/v1/` ends up nowhere near `/chat/completions`,
// so the `/v1/chat/completions` hint misses. `/chat/completions`
// is OpenAI's API contract: any service accepting OpenAI bodies
// serves at this path, so false-positive risk is negligible.
// `/completions` (legacy), `/embeddings`, and `/responses` are
// kept on the canonical-only path because their bare forms are
// too generic to be safe substrings.
var openAIPathHints = []string{
"/v1/chat/completions",
"/v1/completions",
"/v1/embeddings",
"/v1/responses",
"/chat/completions",
}
// Provider returns ProviderOpenAI.
func (OpenAIParser) Provider() Provider { return ProviderOpenAI }
// ProviderName returns the stable label used for metrics and metadata.
func (OpenAIParser) ProviderName() string { return "openai" }
// DetectFromURL reports whether the given request path looks like an OpenAI
// API endpoint. The match is case-insensitive and substring-based so that a
// reverse proxy prefix strip or rewrite does not defeat detection.
func (OpenAIParser) DetectFromURL(path string) bool {
lower := strings.ToLower(path)
for _, hint := range openAIPathHints {
if strings.Contains(lower, hint) {
return true
}
}
return false
}
type openAIRequest struct {
Model string `json:"model"`
Stream *bool `json:"stream"`
StreamOptions *struct {
IncludeUsage *bool `json:"include_usage"`
} `json:"stream_options"`
// Chat Completions / Completions: messages[].content (string or array of
// content parts). Responses API: input is either a string or an array of
// items with content parts. We use json.RawMessage to defer parsing each
// shape independently.
Messages []openAIMessage `json:"messages"`
Prompt json.RawMessage `json:"prompt"`
Input json.RawMessage `json:"input"`
}
type openAIMessage struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
}
// ParseRequest extracts the model name and streaming flag from an OpenAI
// request body. Unknown or missing fields leave the corresponding struct
// members zero-valued.
func (OpenAIParser) ParseRequest(body []byte) (RequestFacts, error) {
var req openAIRequest
if err := json.Unmarshal(body, &req); err != nil {
return RequestFacts{}, fmt.Errorf("decode openai request: %w: %v", ErrMalformedRequest, err)
}
return RequestFacts{
Model: req.Model,
Stream: ptrDeref(req.Stream),
}, nil
}
// openAIResponse accepts both naming conventions in a single struct because
// OpenAI's older Chat Completions API uses prompt_tokens/completion_tokens
// while the newer Responses API (/v1/responses) uses input_tokens/output_tokens
// (aligned with Anthropic). Pointer fields let us tell "absent" from "zero".
//
// PromptTokensDetails.CachedTokens (Chat Completions) and
// InputTokensDetails.CachedTokens (Responses API) carry the SUBSET of
// prompt/input tokens that hit the prompt cache. Cost-meter applies the
// discount rate to that subset and the regular rate to the remainder so
// we never double-bill the cached portion.
type openAIResponse struct {
Usage struct {
PromptTokens *int64 `json:"prompt_tokens"`
CompletionTokens *int64 `json:"completion_tokens"`
InputTokens *int64 `json:"input_tokens"`
OutputTokens *int64 `json:"output_tokens"`
TotalTokens *int64 `json:"total_tokens"`
PromptTokensDetails *struct {
CachedTokens *int64 `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
InputTokensDetails *struct {
CachedTokens *int64 `json:"cached_tokens"`
} `json:"input_tokens_details"`
} `json:"usage"`
}
// ParseResponse decodes the non-streaming OpenAI response envelope. Status
// codes other than 200 are treated as non-LLM responses so the caller can
// skip cost accounting without aborting the request.
func (OpenAIParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) {
if status != 200 {
return Usage{}, fmt.Errorf("openai status %d: %w", status, ErrNotLLMResponse)
}
if isEventStream(contentType) {
return Usage{}, ErrStreamingUnsupported
}
if !isJSON(contentType) {
return Usage{}, fmt.Errorf("openai content-type %q: %w", contentType, ErrNotLLMResponse)
}
var resp openAIResponse
if err := json.Unmarshal(body, &resp); err != nil {
return Usage{}, fmt.Errorf("decode openai response: %w: %v", ErrMalformedResponse, err)
}
// Responses-API names take precedence when present; fall back to the older
// Chat Completions names. This handles both endpoints transparently
// without forcing a per-route configuration.
u := Usage{
InputTokens: pickInt64(resp.Usage.InputTokens, resp.Usage.PromptTokens),
OutputTokens: pickInt64(resp.Usage.OutputTokens, resp.Usage.CompletionTokens),
TotalTokens: derefInt64(resp.Usage.TotalTokens),
CachedInputTokens: openAICachedTokens(resp),
}
if u.TotalTokens == 0 && (u.InputTokens > 0 || u.OutputTokens > 0) {
u.TotalTokens = u.InputTokens + u.OutputTokens
}
return u, nil
}
// openAICachedTokens returns the cached-prompt subset reported by
// either the Responses-API (input_tokens_details.cached_tokens) or
// the Chat-Completions API (prompt_tokens_details.cached_tokens).
// Responses-API takes precedence when both are populated.
func openAICachedTokens(resp openAIResponse) int64 {
// Responses-API details are authoritative when present: an explicit
// cached_tokens of 0 must be honored, not treated as missing and
// overridden by the Chat-Completions field (which would overstate cache).
if resp.Usage.InputTokensDetails != nil && resp.Usage.InputTokensDetails.CachedTokens != nil {
return derefInt64(resp.Usage.InputTokensDetails.CachedTokens)
}
if resp.Usage.PromptTokensDetails != nil {
return derefInt64(resp.Usage.PromptTokensDetails.CachedTokens)
}
return 0
}
// ExtractPrompt returns the user-visible prompt text from an OpenAI request.
// Handles chat.completions (messages[].content), legacy completions (prompt
// string), and the Responses API (input as string or content-part array).
// Returns "" when nothing extractable is found.
func (OpenAIParser) ExtractPrompt(body []byte) string {
var req openAIRequest
if err := json.Unmarshal(body, &req); err != nil {
return ""
}
if len(req.Messages) > 0 {
return joinMessages(req.Messages)
}
if len(req.Input) > 0 {
return extractResponsesInput(req.Input)
}
if len(req.Prompt) > 0 {
return decodeStringOrJoin(req.Prompt)
}
return ""
}
// extractResponsesInput handles the Responses API `input` field. It is one
// of three shapes: a plain string, an array of message items
// ({role, content: string | [parts]}) as sent by Codex and the Responses
// SDK, or a flat array of content parts ({type, text/input_text}). Message
// items are flattened to "role: text" lines; items without extractable text
// (reasoning blocks, tool calls) are skipped.
func extractResponsesInput(raw json.RawMessage) string {
if s, ok := tryDecodeString(raw); ok {
return s
}
var items []struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
Text string `json:"text"`
InputText string `json:"input_text"`
}
if err := json.Unmarshal(raw, &items); err != nil {
return extractContentParts(raw)
}
var b strings.Builder
for _, it := range items {
var text string
switch {
case len(it.Content) > 0:
text = decodeStringOrJoin(it.Content)
case it.Text != "":
text = it.Text
case it.InputText != "":
text = it.InputText
}
if text == "" {
continue
}
if b.Len() > 0 {
b.WriteByte('\n')
}
if it.Role != "" {
b.WriteString(it.Role)
b.WriteString(": ")
}
b.WriteString(text)
}
return b.String()
}
// ExtractSessionID reads the OpenAI session marker. Codex (the Responses
// API client) stamps client_metadata.session_id on every request body;
// plain chat.completions traffic carries no session id and yields "".
func (OpenAIParser) ExtractSessionID(body []byte) string {
var req struct {
ClientMetadata struct {
SessionID string `json:"session_id"`
} `json:"client_metadata"`
}
if err := json.Unmarshal(body, &req); err != nil {
return ""
}
return req.ClientMetadata.SessionID
}
type openAIChatChoice struct {
Message struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
} `json:"message"`
Text string `json:"text"`
}
type openAIChatResponse struct {
Choices []openAIChatChoice `json:"choices"`
// Responses API: output[].content[].text
Output []struct {
Type string `json:"type"`
Content json.RawMessage `json:"content"`
Text string `json:"text"`
} `json:"output"`
OutputText string `json:"output_text"`
}
// ExtractCompletion returns the assistant text from a non-streaming OpenAI
// response. Handles chat.completions (choices[].message.content), legacy
// completions (choices[].text), and Responses API (output[].content[].text
// or the convenience output_text field).
func (OpenAIParser) ExtractCompletion(status int, contentType string, body []byte) string {
if status != 200 || isEventStream(contentType) || !isJSON(contentType) {
return ""
}
var resp openAIChatResponse
if err := json.Unmarshal(body, &resp); err != nil {
return ""
}
if resp.OutputText != "" {
return resp.OutputText
}
for _, c := range resp.Choices {
if len(c.Message.Content) > 0 {
if s := decodeStringOrJoin(c.Message.Content); s != "" {
return s
}
}
if c.Text != "" {
return c.Text
}
}
for _, o := range resp.Output {
if o.Text != "" {
return o.Text
}
if len(o.Content) > 0 {
if s := extractContentParts(o.Content); s != "" {
return s
}
}
}
return ""
}
// joinMessages flattens a chat.completions messages array into a single
// "role: content" string per message, separated by newlines. Roles surface
// system/user/assistant context which is useful for log review.
func joinMessages(msgs []openAIMessage) string {
var b strings.Builder
for i, m := range msgs {
if i > 0 {
b.WriteByte('\n')
}
if m.Role != "" {
b.WriteString(m.Role)
b.WriteString(": ")
}
b.WriteString(decodeStringOrJoin(m.Content))
}
return b.String()
}
// extractContentParts handles the Responses-API content shape, which is
// either a single string or an array of {type, text} parts. text and
// input_text both carry user-facing content.
func extractContentParts(raw json.RawMessage) string {
if s, ok := tryDecodeString(raw); ok {
return s
}
var parts []struct {
Type string `json:"type"`
Text string `json:"text"`
InputText string `json:"input_text"`
}
if err := json.Unmarshal(raw, &parts); err != nil {
// Last-ditch: array of strings.
var arr []string
if json.Unmarshal(raw, &arr) == nil {
return strings.Join(arr, "\n")
}
return ""
}
var b strings.Builder
for _, p := range parts {
var text string
switch {
case p.Text != "":
text = p.Text
case p.InputText != "":
text = p.InputText
}
if text == "" {
continue
}
if b.Len() > 0 {
b.WriteByte('\n')
}
b.WriteString(text)
}
return b.String()
}
// decodeStringOrJoin accepts either a JSON string or a content-parts array
// (chat.completions multimodal) and returns a flat string. Multimodal parts
// are separated by newlines; non-text parts are skipped.
func decodeStringOrJoin(raw json.RawMessage) string {
if s, ok := tryDecodeString(raw); ok {
return s
}
return extractContentParts(raw)
}
func tryDecodeString(raw json.RawMessage) (string, bool) {
if len(raw) == 0 {
return "", false
}
var s string
if err := json.Unmarshal(raw, &s); err == nil {
return s, true
}
return "", false
}
// pickInt64 returns the first non-nil pointer's value. Used to prefer one
// naming convention while transparently falling back to another.
func pickInt64(preferred, fallback *int64) int64 {
if preferred != nil {
return *preferred
}
return derefInt64(fallback)
}
func derefInt64(v *int64) int64 {
if v == nil {
return 0
}
return *v
}
func ptrDeref(b *bool) bool {
if b == nil {
return false
}
return *b
}
func isEventStream(contentType string) bool {
return strings.Contains(strings.ToLower(contentType), "text/event-stream")
}
func isJSON(contentType string) bool {
lower := strings.ToLower(contentType)
return strings.Contains(lower, "application/json") || strings.Contains(lower, "+json")
}
+255
View File
@@ -0,0 +1,255 @@
package llm
import (
"errors"
"os"
"path/filepath"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestOpenAIDetectFromURL(t *testing.T) {
p := OpenAIParser{}
cases := map[string]bool{
"/v1/chat/completions": true,
"/v1/completions": true,
"/v1/embeddings": true,
"/v1/responses": true,
"/API/V1/Chat/Completions": true,
"/upstream/v1/chat/completions?trace=1": true,
// Cloudflare AI Gateway puts its own /v1/{account}/{gateway}
// segment between the canonical /v1/ and the provider slug,
// so the /v1/chat/completions substring no longer appears
// adjacent in the path. The bare /chat/completions hint
// catches Cloudflare's OpenAI direct path
// (/v1/{account}/{gateway}/openai/chat/completions) and
// compat path (/v1/{account}/{gateway}/compat/chat/completions).
"/v1/{account}/{gateway}/openai/chat/completions": true,
"/v1/{account}/{gateway}/compat/chat/completions": true,
"/chat/completions": true,
"/v1/messages": false,
"/healthz": false,
"": false,
}
for path, want := range cases {
assert.Equal(t, want, p.DetectFromURL(path), "DetectFromURL(%q)", path)
}
}
func TestOpenAIParseRequest(t *testing.T) {
p := OpenAIParser{}
t.Run("stream true", func(t *testing.T) {
facts, err := p.ParseRequest([]byte(`{"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true}}`))
require.NoError(t, err)
assert.Equal(t, "gpt-4o", facts.Model, "request model extracted")
assert.True(t, facts.Stream, "request marked as streaming")
})
t.Run("stream default", func(t *testing.T) {
facts, err := p.ParseRequest([]byte(`{"model":"gpt-4o-mini"}`))
require.NoError(t, err)
assert.Equal(t, "gpt-4o-mini", facts.Model, "request model extracted")
assert.False(t, facts.Stream, "missing stream flag defaults to false")
})
t.Run("malformed", func(t *testing.T) {
_, err := p.ParseRequest([]byte(`{not json}`))
require.Error(t, err)
assert.True(t, errors.Is(err, ErrMalformedRequest), "sentinel error wrapped")
})
}
func TestOpenAIParseResponse(t *testing.T) {
p := OpenAIParser{}
t.Run("happy fixture", func(t *testing.T) {
body, err := os.ReadFile(filepath.Join("fixtures", "openai_chat_completion.json"))
require.NoError(t, err, "fixture must be readable")
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(123), usage.InputTokens, "prompt tokens become input")
assert.Equal(t, int64(45), usage.OutputTokens, "completion tokens become output")
assert.Equal(t, int64(168), usage.TotalTokens, "total_tokens carried through")
})
t.Run("total computed when missing", func(t *testing.T) {
body := []byte(`{"usage":{"prompt_tokens":10,"completion_tokens":5}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(15), usage.TotalTokens, "total computed from in+out")
})
t.Run("streaming rejected", func(t *testing.T) {
_, err := p.ParseResponse(200, "text/event-stream", []byte(""))
require.ErrorIs(t, err, ErrStreamingUnsupported, "SSE responses must use the scanner")
})
t.Run("non-200", func(t *testing.T) {
_, err := p.ParseResponse(500, "application/json", []byte(`{"error":"x"}`))
require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 rejected as non-LLM")
})
t.Run("non-json content type", func(t *testing.T) {
_, err := p.ParseResponse(200, "text/plain", []byte(`{}`))
require.ErrorIs(t, err, ErrNotLLMResponse, "text/plain treated as non-LLM")
})
t.Run("malformed body", func(t *testing.T) {
_, err := p.ParseResponse(200, "application/json", []byte(`{not json`))
require.ErrorIs(t, err, ErrMalformedResponse, "bad JSON yields malformed error")
})
// Responses-API fixture: /v1/responses returns input_tokens/output_tokens
// (Anthropic-style) instead of prompt_tokens/completion_tokens. The parser
// must accept both.
t.Run("responses api fixture", func(t *testing.T) {
body, err := os.ReadFile(filepath.Join("fixtures", "openai_responses.json"))
require.NoError(t, err, "fixture must be readable")
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(15), usage.InputTokens, "input_tokens should map directly")
assert.Equal(t, int64(414), usage.OutputTokens, "output_tokens should map directly")
assert.Equal(t, int64(429), usage.TotalTokens, "total_tokens carried through")
})
t.Run("responses api naming preferred over chat-completions when both present", func(t *testing.T) {
body := []byte(`{"usage":{"prompt_tokens":1,"completion_tokens":2,"input_tokens":15,"output_tokens":414,"total_tokens":429}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(15), usage.InputTokens, "responses-api names take precedence")
assert.Equal(t, int64(414), usage.OutputTokens, "responses-api names take precedence")
})
t.Run("chat-completions naming still works alone", func(t *testing.T) {
body := []byte(`{"usage":{"prompt_tokens":15,"completion_tokens":414,"total_tokens":429}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(15), usage.InputTokens, "prompt_tokens fallback")
assert.Equal(t, int64(414), usage.OutputTokens, "completion_tokens fallback")
})
// Cached-prompt accounting. cached_tokens is a SUBSET of
// prompt_tokens — input_tokens carries the full prompt count and
// the cached subset is reported separately so the cost meter can
// apply the discount rate to that portion.
t.Run("chat-completions cached_tokens subset surfaces", func(t *testing.T) {
body := []byte(`{"usage":{"prompt_tokens":1024,"completion_tokens":200,"total_tokens":1224,"prompt_tokens_details":{"cached_tokens":768}}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(1024), usage.InputTokens, "input remains the full prompt count — cached is a subset, not a separate bucket")
assert.Equal(t, int64(768), usage.CachedInputTokens, "cached_tokens must surface so cost meter can discount the cached subset")
assert.Zero(t, usage.CacheCreationTokens, "OpenAI has no cache_creation analogue")
})
t.Run("responses-api input_tokens_details.cached_tokens surfaces", func(t *testing.T) {
body := []byte(`{"usage":{"input_tokens":2048,"output_tokens":100,"total_tokens":2148,"input_tokens_details":{"cached_tokens":1500}}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(2048), usage.InputTokens)
assert.Equal(t, int64(1500), usage.CachedInputTokens, "Responses-API input_tokens_details.cached_tokens path must surface too")
})
t.Run("responses-api cached takes precedence over chat-completions when both present", func(t *testing.T) {
body := []byte(`{"usage":{"prompt_tokens":1,"input_tokens":2,"output_tokens":3,"prompt_tokens_details":{"cached_tokens":50},"input_tokens_details":{"cached_tokens":99}}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Equal(t, int64(99), usage.CachedInputTokens, "Responses-API field wins when both naming conventions are present")
})
t.Run("absent cached_tokens leaves cached counts at zero", func(t *testing.T) {
body := []byte(`{"usage":{"prompt_tokens":15,"completion_tokens":414,"total_tokens":429}}`)
usage, err := p.ParseResponse(200, "application/json", body)
require.NoError(t, err)
assert.Zero(t, usage.CachedInputTokens, "no prompt_tokens_details = no cached subset")
})
}
func TestOpenAIExtractPrompt_ChatCompletions(t *testing.T) {
body := []byte(`{"model":"gpt-4o-mini","messages":[{"role":"system","content":"be brief"},{"role":"user","content":"ping"}]}`)
got := OpenAIParser{}.ExtractPrompt(body)
require.NotEmpty(t, got, "messages array must extract")
require.Contains(t, got, "system: be brief", "system role and content surface")
require.Contains(t, got, "user: ping", "user role and content surface")
}
func TestOpenAIExtractPrompt_ResponsesAPIStringInput(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","input":"Hello there"}`)
got := OpenAIParser{}.ExtractPrompt(body)
require.Equal(t, "Hello there", got, "string input field should pass through")
}
func TestOpenAIExtractPrompt_ResponsesAPIInputParts(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","input":[{"type":"input_text","input_text":"first"},{"type":"input_text","input_text":"second"}]}`)
got := OpenAIParser{}.ExtractPrompt(body)
require.Contains(t, got, "first", "first content part surfaces")
require.Contains(t, got, "second", "second content part surfaces")
}
// TestOpenAIExtractPrompt_ResponsesAPIMessageItems guards the live Codex
// shape: input is an array of message items whose text is nested under
// content[].text, not flat content parts. The old code fed the outer array
// to the content-part decoder and extracted nothing, so the stored prompt
// was empty.
func TestOpenAIExtractPrompt_ResponsesAPIMessageItems(t *testing.T) {
body := []byte(`{"model":"gpt-5.5","input":[` +
`{"type":"message","role":"developer","content":[{"type":"input_text","text":"system rules"}]},` +
`{"type":"message","role":"user","content":[{"type":"input_text","text":"hello there"}]},` +
`{"type":"reasoning","encrypted_content":"opaque","summary":[]},` +
`{"type":"message","role":"assistant","content":[{"type":"output_text","text":"prior reply"}]}` +
`]}`)
got := OpenAIParser{}.ExtractPrompt(body)
require.Contains(t, got, "system rules", "developer message content must surface")
require.Contains(t, got, "hello there", "user message content must surface")
require.Contains(t, got, "developer:", "role labels must prefix each message")
require.NotContains(t, got, "opaque", "reasoning items without text must be skipped")
}
func TestOpenAIExtractPrompt_LegacyCompletion(t *testing.T) {
body := []byte(`{"model":"text-davinci-003","prompt":"once upon a time"}`)
got := OpenAIParser{}.ExtractPrompt(body)
require.Equal(t, "once upon a time", got, "string prompt field should pass through")
}
func TestOpenAIExtractSessionID(t *testing.T) {
t.Run("codex client_metadata.session_id", func(t *testing.T) {
body := []byte(`{"model":"gpt-5.5","client_metadata":{"session_id":"019eeb72-ab7c-7cd2","thread_id":"t1"},"input":[]}`)
assert.Equal(t, "019eeb72-ab7c-7cd2", OpenAIParser{}.ExtractSessionID(body), "Codex session id must come from client_metadata.session_id")
})
t.Run("plain chat has no session", func(t *testing.T) {
body := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
assert.Equal(t, "", OpenAIParser{}.ExtractSessionID(body), "plain chat.completions carries no session id")
})
t.Run("non-JSON yields empty", func(t *testing.T) {
assert.Equal(t, "", OpenAIParser{}.ExtractSessionID([]byte("not json")), "malformed body must not error")
})
}
func TestOpenAIExtractCompletion_ChatCompletions(t *testing.T) {
body, err := os.ReadFile(filepath.Join("fixtures", "openai_chat_completion.json"))
require.NoError(t, err)
got := OpenAIParser{}.ExtractCompletion(200, "application/json", body)
require.NotEmpty(t, got, "fixture has assistant content")
}
func TestOpenAIExtractCompletion_ResponsesAPI(t *testing.T) {
body, err := os.ReadFile(filepath.Join("fixtures", "openai_responses.json"))
require.NoError(t, err)
got := OpenAIParser{}.ExtractCompletion(200, "application/json", body)
require.NotEmpty(t, got, "responses-api fixture has output content")
}
func TestOpenAIExtractCompletion_Streaming(t *testing.T) {
got := OpenAIParser{}.ExtractCompletion(200, "text/event-stream", []byte(""))
require.Empty(t, got, "streaming responses are skipped")
}
func TestOpenAIExtractCompletion_NonOK(t *testing.T) {
got := OpenAIParser{}.ExtractCompletion(500, "application/json", []byte(`{"choices":[{"message":{"content":"x"}}]}`))
require.Empty(t, got, "non-200 returns empty")
}
+112
View File
@@ -0,0 +1,112 @@
// Package llm provides the shared LLM request and response parsing
// library consumed by proxy middleware. It is runtime agnostic: the same
// package is used by the native built-in executor now and will be reused
// by the WASM adapter later.
package llm
// Provider identifies an LLM API provider.
type Provider int
const (
// ProviderUnknown signals that no parser matched the request.
ProviderUnknown Provider = 0
// ProviderOpenAI identifies the OpenAI API surface.
ProviderOpenAI Provider = 1
// ProviderAnthropic identifies the Anthropic Messages API surface.
ProviderAnthropic Provider = 2
// ProviderBedrock identifies the AWS Bedrock runtime surface.
ProviderBedrock Provider = 3
)
// RequestFacts captures the subset of the LLM request body that the
// middleware annotates as metadata (model, streaming flag). Additional
// fields are added as parsers grow.
type RequestFacts struct {
Model string
Stream bool
}
// Usage is the provider-agnostic token accounting emitted to metrics and
// access logs. Downstream consumers map InputTokens/OutputTokens to the
// plg.llm.* metadata allowlist entries.
//
// CachedInputTokens carries OpenAI's prompt_tokens_details.cached_tokens
// (a SUBSET of InputTokens) when the response is from OpenAI, or
// Anthropic's cache_read_input_tokens (ADDITIVE to InputTokens) when from
// Anthropic. The cost meter switches formula on KeyLLMProvider so the
// two shapes are billed correctly without double-counting.
//
// CacheCreationTokens carries Anthropic's cache_creation_input_tokens
// (ADDITIVE; not present in the OpenAI shape).
type Usage struct {
InputTokens int64
OutputTokens int64
TotalTokens int64
CachedInputTokens int64
CacheCreationTokens int64
}
// Parser is the per-provider interface implemented in this package. The
// dispatcher selects a parser by calling DetectFromURL against the incoming
// request path; ties break by registration order (see Parsers).
type Parser interface {
Provider() Provider
ProviderName() string
DetectFromURL(path string) bool
ParseRequest(body []byte) (RequestFacts, error)
ParseResponse(status int, contentType string, body []byte) (Usage, error)
// ExtractPrompt returns the user-facing prompt text from a request body.
// Different endpoint shapes (chat.completions, responses, messages) are
// handled by the per-provider implementation. Returns "" when no prompt
// can be extracted; never returns an error — extraction is best-effort
// because callers use the result for observability, not authorization.
ExtractPrompt(body []byte) string
// ExtractCompletion returns the assistant-facing completion text from a
// non-streaming response body. status and contentType match the
// ParseResponse arguments so implementations can fast-fail uniformly.
ExtractCompletion(status int, contentType string, body []byte) string
// ExtractSessionID returns a stable identifier that groups requests of
// the same conversation / coding session, read from the per-provider
// location clients populate (e.g. OpenAI Codex's client_metadata.session_id,
// Claude Code's metadata.user_id). Returns "" when the body carries no
// recognised session marker; extraction is best-effort and never errors.
ExtractSessionID(body []byte) string
}
// Parsers returns the built-in parser set in a stable order. The order is
// deterministic so that DetectFromURL ties produce consistent routing.
func Parsers() []Parser {
return []Parser{
OpenAIParser{},
AnthropicParser{},
BedrockParser{},
}
}
// DetectParser returns the first parser whose DetectFromURL matches the given
// request path. ok=false means no parser claimed the path.
func DetectParser(path string) (Parser, bool) {
for _, p := range Parsers() {
if p.DetectFromURL(path) {
return p, true
}
}
return nil, false
}
// ParserByName returns the parser whose ProviderName matches id. Used by
// callers that already know which provider surface a request will hit
// (e.g. the agent-network middleware chain configured per synthesised
// service) so they can skip URL sniffing. ok=false when no parser is
// registered under that name.
func ParserByName(id string) (Parser, bool) {
if id == "" {
return nil, false
}
for _, p := range Parsers() {
if p.ProviderName() == id {
return p, true
}
}
return nil, false
}
+54
View File
@@ -0,0 +1,54 @@
package llm
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestParsers_ProviderNames(t *testing.T) {
parsers := Parsers()
require.Len(t, parsers, 3, "three built-in parsers expected")
names := make([]string, 0, len(parsers))
for _, p := range parsers {
names = append(names, p.ProviderName())
}
assert.Contains(t, names, "openai", "OpenAI parser should be registered")
assert.Contains(t, names, "anthropic", "Anthropic parser should be registered")
assert.Contains(t, names, "bedrock", "Bedrock parser should be registered")
}
func TestDetectParser(t *testing.T) {
cases := []struct {
name string
path string
expectedName string
expectOK bool
}{
{"openai chat", "/v1/chat/completions", "openai", true},
{"openai prefixed", "/api/v1/chat/completions", "openai", true},
{"openai responses", "/v1/responses", "openai", true},
{"anthropic messages", "/v1/messages", "anthropic", true},
{"anthropic prefixed", "/proxy/v1/messages?query", "anthropic", true},
{"unknown path", "/healthz", "", false},
{"empty path", "", "", false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
p, ok := DetectParser(tc.path)
require.Equal(t, tc.expectOK, ok, "detection success mismatch for %q", tc.path)
if ok {
assert.Equal(t, tc.expectedName, p.ProviderName(), "provider name mismatch")
}
})
}
}
func TestProviderValues(t *testing.T) {
assert.Equal(t, Provider(0), ProviderUnknown, "unknown provider is the zero value")
assert.Equal(t, ProviderOpenAI, OpenAIParser{}.Provider(), "OpenAI parser reports its provider enum")
assert.Equal(t, ProviderAnthropic, AnthropicParser{}.Provider(), "Anthropic parser reports its provider enum")
}
@@ -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
+449
View File
@@ -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)
}
+432
View File
@@ -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
}
+117
View File
@@ -0,0 +1,117 @@
package llm
import (
"bufio"
"errors"
"fmt"
"io"
"strings"
)
// Event represents a single server-sent event. Type is the dispatch name
// carried on an "event:" line (empty when the stream uses only "data:"
// lines). Data is the concatenation of every "data:" line that made up the
// event, joined by a single newline.
type Event struct {
Type string
Data string
}
// Scanner reads SSE events from an underlying byte stream. Events are
// delimited by a blank line ("\n\n"). CRLF line endings are normalized to LF
// transparently so fixtures captured from live servers can be replayed.
//
// Scanner is not safe for concurrent use.
type Scanner struct {
r *bufio.Reader
maxLine int
}
// NewScanner wraps the given reader. The default underlying buffer size is
// large enough for typical provider events (~64 KiB); callers needing
// larger events can wrap the reader in their own bufio.Reader beforehand.
func NewScanner(r io.Reader) *Scanner {
return &Scanner{
r: bufio.NewReaderSize(r, 64*1024),
maxLine: 1 << 20,
}
}
// Next returns the next event. It returns io.EOF after the final event has
// been consumed. A trailing event that is not terminated by a blank line is
// still returned before io.EOF so that servers which close the connection
// without a trailing newline are handled correctly.
func (s *Scanner) Next() (Event, error) {
var (
event Event
dataBuf strings.Builder
hasData bool
hasAny bool
)
for {
line, err := s.readLine()
if err != nil {
if errors.Is(err, io.EOF) && hasAny {
event.Data = dataBuf.String()
return event, nil
}
return Event{}, err
}
if line == "" {
if !hasAny {
continue
}
event.Data = dataBuf.String()
return event, nil
}
hasAny = true
if strings.HasPrefix(line, ":") {
continue
}
field, value := splitField(line)
switch field {
case "event":
event.Type = value
case "data":
if hasData {
dataBuf.WriteByte('\n')
}
dataBuf.WriteString(value)
hasData = true
}
}
}
func (s *Scanner) readLine() (string, error) {
line, err := s.r.ReadString('\n')
if err != nil {
if errors.Is(err, io.EOF) && line != "" {
return trimEOL(line), nil
}
return "", err
}
if len(line) > s.maxLine {
return "", fmt.Errorf("sse line exceeds %d bytes", s.maxLine)
}
return trimEOL(line), nil
}
func trimEOL(line string) string {
line = strings.TrimRight(line, "\n")
line = strings.TrimRight(line, "\r")
return line
}
func splitField(line string) (string, string) {
idx := strings.IndexByte(line, ':')
if idx < 0 {
return line, ""
}
field := line[:idx]
value := strings.TrimPrefix(line[idx+1:], " ")
return field, value
}
+175
View File
@@ -0,0 +1,175 @@
package llm
import (
"errors"
"io"
"os"
"path/filepath"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func collectEvents(t *testing.T, r io.Reader) []Event {
t.Helper()
s := NewScanner(r)
var out []Event
for {
ev, err := s.Next()
if errors.Is(err, io.EOF) {
return out
}
require.NoError(t, err, "unexpected error scanning SSE")
out = append(out, ev)
}
}
func TestSSEScanner_OpenAIFixture(t *testing.T) {
f, err := os.Open(filepath.Join("fixtures", "openai_stream.txt"))
require.NoError(t, err, "fixture must be openable")
defer f.Close()
events := collectEvents(t, f)
require.Len(t, events, 4, "expected 4 data frames (3 chunks + [DONE])")
for _, ev := range events {
assert.Empty(t, ev.Type, "OpenAI stream uses data-only frames")
}
assert.Contains(t, events[2].Data, `"usage"`, "third chunk carries usage block")
assert.Equal(t, "[DONE]", events[3].Data, "final frame is the OpenAI DONE sentinel")
}
func TestSSEScanner_AnthropicFixture(t *testing.T) {
f, err := os.Open(filepath.Join("fixtures", "anthropic_stream.txt"))
require.NoError(t, err, "fixture must be openable")
defer f.Close()
events := collectEvents(t, f)
require.Len(t, events, 7, "expected 7 Anthropic events")
types := make([]string, 0, len(events))
for _, ev := range events {
types = append(types, ev.Type)
}
assert.Equal(t, []string{
"message_start",
"content_block_start",
"content_block_delta",
"content_block_delta",
"content_block_stop",
"message_delta",
"message_stop",
}, types, "Anthropic event ordering matches fixture")
var deltaUsage Event
for _, ev := range events {
if ev.Type == "message_delta" {
deltaUsage = ev
break
}
}
assert.Contains(t, deltaUsage.Data, `"output_tokens":45`, "message_delta carries partial usage")
}
func TestSSEScanner_MultilineData(t *testing.T) {
raw := "event: ping\ndata: line1\ndata: line2\ndata: line3\n\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 1, "one logical event from three data lines")
assert.Equal(t, "ping", events[0].Type, "event name honored")
assert.Equal(t, "line1\nline2\nline3", events[0].Data, "data lines joined with newline")
}
func TestSSEScanner_CRLF(t *testing.T) {
raw := "event: foo\r\ndata: bar\r\n\r\ndata: baz\r\n\r\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 2, "CRLF-delimited events recognized")
assert.Equal(t, "foo", events[0].Type, "first event type preserved")
assert.Equal(t, "bar", events[0].Data, "first event data preserved")
assert.Empty(t, events[1].Type, "second event has no event name")
assert.Equal(t, "baz", events[1].Data, "second event data preserved")
}
func TestSSEScanner_EmptyInput(t *testing.T) {
s := NewScanner(strings.NewReader(""))
_, err := s.Next()
require.ErrorIs(t, err, io.EOF, "empty input yields immediate EOF")
}
func TestSSEScanner_CommentIgnored(t *testing.T) {
raw := ": this is a comment\ndata: hi\n\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 1, "comment line does not emit an event")
assert.Equal(t, "hi", events[0].Data, "data line honoured after comment")
}
func TestSSEScanner_TrailingWithoutBlankLine(t *testing.T) {
raw := "event: foo\ndata: bar\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 1, "trailing event without blank line still emitted")
assert.Equal(t, "foo", events[0].Type)
assert.Equal(t, "bar", events[0].Data)
}
// TestSSEScanner_ManyConsecutiveEmptyLines feeds a stream that is nothing
// but empty lines. The scanner must terminate without panic — empty lines
// alone do not constitute an event and must yield io.EOF.
func TestSSEScanner_ManyConsecutiveEmptyLines(t *testing.T) {
raw := strings.Repeat("\n", 100)
s := NewScanner(strings.NewReader(raw))
_, err := s.Next()
require.ErrorIs(t, err, io.EOF, "100 empty lines must terminate as EOF without panic")
}
// TestSSEScanner_InterleavedCRLFAndLF mixes \r\n and \n terminators within
// the same event. The scanner normalizes both and must still recover a
// coherent event.
func TestSSEScanner_InterleavedCRLFAndLF(t *testing.T) {
raw := "event: mix\r\ndata: first\ndata: second\r\n\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 1, "mixed line endings must still produce one event")
assert.Equal(t, "mix", events[0].Type)
assert.Equal(t, "first\nsecond", events[0].Data, "both data lines joined")
}
// TestSSEScanner_LongSingleDataLine constructs a single data line that
// exceeds the default bufio buffer (64 KiB) but stays under the scanner
// maxLine. The scanner must round-trip the value intact without panicking
// or truncating silently.
func TestSSEScanner_LongSingleDataLine(t *testing.T) {
big := strings.Repeat("x", 80<<10)
raw := "data: " + big + "\n\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 1, "long single-line event must be emitted")
assert.Equal(t, big, events[0].Data, "long data preserved")
}
// TestSSEScanner_BinaryGarbageInData validates that non-printable bytes
// inside a data line do not crash the parser. The scanner should either
// round-trip them or return a well-formed error — never panic.
func TestSSEScanner_BinaryGarbageInData(t *testing.T) {
raw := "data: \x00\x01\x02\xff\xfe\n\n"
defer func() {
if r := recover(); r != nil {
t.Fatalf("scanner panicked on binary garbage: %v", r)
}
}()
s := NewScanner(strings.NewReader(raw))
ev, err := s.Next()
require.NoError(t, err, "binary bytes in data should not surface as error")
assert.Equal(t, "\x00\x01\x02\xff\xfe", ev.Data, "binary payload round-trips")
}
// TestSSEScanner_UnknownFieldsIgnored stresses the field parser by sending
// unrecognized field names ("id:", "retry:", "custom:"). They must be
// silently ignored per the SSE spec; the scanner must not panic or emit
// spurious events.
func TestSSEScanner_UnknownFieldsIgnored(t *testing.T) {
raw := "id: 1\nretry: 5000\ncustom: value\ndata: payload\n\n"
events := collectEvents(t, strings.NewReader(raw))
require.Len(t, events, 1, "unknown fields must not spawn extra events")
assert.Equal(t, "payload", events[0].Data, "data field survives amid unknown fields")
}