From e6ac1ad9b525ff45263ca5f0d92129bca397fac5 Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Mon, 3 Aug 2026 15:10:27 +1000 Subject: [PATCH] feat: Ochre usage/cost metering, token quota, and RPM throttle (#675) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the .7.6 metering slice on top of .7.4's invocation stream: a second durable consumer (ochre-usage-metering) aggregates per-account token counts and external cost into a new bedrock-usage KV bucket, keyed {accountID}/{modelID}/{period}. RequestID dedupe (an atomic create-if-absent marker in bedrock-usage-dedupe) makes redelivery count-neutral rather than double-billing a customer. Prices are integer micro-USD per million tokens (D12): in-tree catalog defaults with KV overrides in bedrock-prices, mirroring the serving-spec shape from .7.3. Self-hosted models always resolve a known zero; a model with neither an override nor a default resolves Known=false, which the usage counter surfaces as a sticky CostUnknown flag distinct from a genuine zero-cost period. handlers_quota gains two Bedrock dimensions, both defaulting to disabled independently of the block's existing Enabled toggle: tokens-per-month (a check-against-history hard cap reading the stream-fed counter, tolerant of a few seconds of lag) and requests-per-minute (a local per-account token bucket enforced immediately in-process, with a periodic KV sync for cross-node observability only — enforcement never round-trips). Both are wired into BedrockRuntime_Request ahead of every Converse/InvokeModel variant. --- spinifex/awserrors/awserrors.go | 17 +- spinifex/gateway/bedrock/catalog.go | 40 ++- spinifex/gateway/bedrock/prices.go | 159 +++++++++ spinifex/gateway/bedrock/prices_test.go | 123 +++++++ spinifex/gateway/bedrock/usage.go | 440 ++++++++++++++++++++++++ spinifex/gateway/bedrock/usage_test.go | 290 ++++++++++++++++ spinifex/gateway/bedrockruntime.go | 11 + spinifex/handlers/quota/bedrock.go | 257 ++++++++++++++ spinifex/handlers/quota/bedrock_test.go | 232 +++++++++++++ spinifex/handlers/quota/quota.go | 26 +- spinifex/services/awsgw/awsgw.go | 23 ++ 11 files changed, 1599 insertions(+), 19 deletions(-) create mode 100644 spinifex/gateway/bedrock/prices.go create mode 100644 spinifex/gateway/bedrock/prices_test.go create mode 100644 spinifex/gateway/bedrock/usage.go create mode 100644 spinifex/gateway/bedrock/usage_test.go create mode 100644 spinifex/handlers/quota/bedrock.go create mode 100644 spinifex/handlers/quota/bedrock_test.go diff --git a/spinifex/awserrors/awserrors.go b/spinifex/awserrors/awserrors.go index 11da4fceb..1c435555e 100644 --- a/spinifex/awserrors/awserrors.go +++ b/spinifex/awserrors/awserrors.go @@ -546,6 +546,10 @@ var ( ErrorModelNotReadyException = "ModelNotReadyException" ErrorServiceUnavailableException = "ServiceUnavailableException" ErrorModelErrorException = "ModelErrorException" + // ErrorServiceQuotaExceededException is for genuine per-account limits + // (e.g. Ochre's tokens-per-month cap) — unlike ErrorModelNotReadyException, + // which is reserved for transient capacity conditions that breach no quota. + ErrorServiceQuotaExceededException = "ServiceQuotaExceededException" ) // ValidErrorCode returns the error code if it exists in ErrorLookup, @@ -1185,10 +1189,11 @@ var ErrorLookup = map[string]ErrorMessage{ // Bedrock/bedrock-runtime error codes. ResourceNotFoundException reuses the // EKS entry (ErrorEKSResourceNotFound) — same wire code, same HTTP status. - ErrorValidationException: {HTTPCode: 400, Message: "The input fails to satisfy the constraints specified by the model or service."}, - ErrorAccessDeniedException: {HTTPCode: 403, Message: "You do not have sufficient access to perform this action."}, - ErrorThrottlingException: {HTTPCode: 429, Message: "The request was denied due to request throttling."}, - ErrorModelNotReadyException: {HTTPCode: 429, Message: "The model specified in the request is not ready to serve inference requests."}, - ErrorServiceUnavailableException: {HTTPCode: 503, Message: "The service isn't currently available. Try again later."}, - ErrorModelErrorException: {HTTPCode: 424, Message: "The request failed because of an error while running the model."}, + ErrorValidationException: {HTTPCode: 400, Message: "The input fails to satisfy the constraints specified by the model or service."}, + ErrorAccessDeniedException: {HTTPCode: 403, Message: "You do not have sufficient access to perform this action."}, + ErrorThrottlingException: {HTTPCode: 429, Message: "The request was denied due to request throttling."}, + ErrorModelNotReadyException: {HTTPCode: 429, Message: "The model specified in the request is not ready to serve inference requests."}, + ErrorServiceUnavailableException: {HTTPCode: 503, Message: "The service isn't currently available. Try again later."}, + ErrorModelErrorException: {HTTPCode: 424, Message: "The request failed because of an error while running the model."}, + ErrorServiceQuotaExceededException: {HTTPCode: 400, Message: "The number of requests exceeds the service quota. Resubmit your request later."}, } diff --git a/spinifex/gateway/bedrock/catalog.go b/spinifex/gateway/bedrock/catalog.go index 02cd40dc0..3f88c947a 100644 --- a/spinifex/gateway/bedrock/catalog.go +++ b/spinifex/gateway/bedrock/catalog.go @@ -28,19 +28,31 @@ const ( // MinVRAMMiB, InstanceType and VLLMArgs are the in-tree serving spec for // self-host entries; the actual staged weights artifact is deployment-local // state resolved from KV instead (see WeightsResolver). +// +// InputPriceMicroUSDPerMillion/OutputPriceMicroUSDPerMillion and PriceKnown +// are the in-tree default price (D12): integer micro-USD per million tokens, +// same shape as the serving spec above. PriceKnown distinguishes an +// explicitly-priced provider entry from one nobody has priced yet — a bare +// zero value is ambiguous with a genuine $0 price, so resolvePrice never +// reads the two int fields without checking it first. Self-host entries leave +// all three at their zero value: resolvePrice special-cases tierSelfHost to +// always return a known zero price rather than consulting these fields. type catalogEntry struct { - ModelID string - ModelName string - ProviderName string - Provider string // "self-host" or "provider:" - InputModalities []string - OutputModalities []string - ResponseStreamingSupported bool - InferenceTypesSupported []string - CustomizationsSupported []string - MinVRAMMiB int // self-host only; 0 for provider entries - InstanceType string // self-host only; system-instance type a launcher boots - VLLMArgs []string // self-host only; extra vLLM server args this model needs + ModelID string + ModelName string + ProviderName string + Provider string // "self-host" or "provider:" + InputModalities []string + OutputModalities []string + ResponseStreamingSupported bool + InferenceTypesSupported []string + CustomizationsSupported []string + MinVRAMMiB int // self-host only; 0 for provider entries + InstanceType string // self-host only; system-instance type a launcher boots + VLLMArgs []string // self-host only; extra vLLM server args this model needs + InputPriceMicroUSDPerMillion int64 // provider entries only; meaningless unless PriceKnown + OutputPriceMicroUSDPerMillion int64 // provider entries only; meaningless unless PriceKnown + PriceKnown bool // provider entries only; self-host is always known-zero } // catalog is the static Phase-1 model set: one self-hosted open model and one @@ -71,6 +83,10 @@ var catalog = []catalogEntry{ OutputModalities: []string{"TEXT"}, ResponseStreamingSupported: false, InferenceTypesSupported: []string{"ON_DEMAND"}, + // List pricing at launch: $3/MTok input, $15/MTok output. + InputPriceMicroUSDPerMillion: 3_000_000, + OutputPriceMicroUSDPerMillion: 15_000_000, + PriceKnown: true, }, } diff --git a/spinifex/gateway/bedrock/prices.go b/spinifex/gateway/bedrock/prices.go new file mode 100644 index 000000000..d6094681f --- /dev/null +++ b/spinifex/gateway/bedrock/prices.go @@ -0,0 +1,159 @@ +package gateway_bedrock + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync" + + "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/nats-io/nats.go/jetstream" +) + +// Price is a resolved per-model price in integer micro-USD per million +// tokens (D12). Known is false when no price could be resolved at all — the +// caller MUST treat that distinctly from a resolved zero price. Self-hosted +// models resolve to a known zero (their real cost is GPU-hours, already +// accounted by the instance); a provider model with neither a KV override +// nor an in-tree default resolves to Known=false, and the two must never be +// mistaken for one another in a billing counter. +type Price struct { + InputMicroUSDPerMillion int64 `json:"inputMicroUsdPerMillion"` + OutputMicroUSDPerMillion int64 `json:"outputMicroUsdPerMillion"` + Known bool `json:"known"` +} + +// PriceResolver resolves modelID's KV-overridden price, if one has been set. +// ok is false when no override exists, in which case resolvePrice falls back +// to the catalog entry's in-tree default. +type PriceResolver interface { + Resolve(ctx context.Context, modelID string) (price Price, ok bool, err error) +} + +// resolvePrice determines entry's price: self-host is always a known zero: a +// KV override or in-tree default is the resolution order for everything +// else, and neither existing leaves the price unknown rather than free. store +// may be nil, in which case only the in-tree default is consulted. +func resolvePrice(ctx context.Context, store PriceResolver, entry catalogEntry) (Price, error) { + if entry.Provider == tierSelfHost { + return Price{Known: true}, nil + } + if store != nil { + price, ok, err := store.Resolve(ctx, entry.ModelID) + if err != nil { + return Price{}, fmt.Errorf("resolve price override for %s: %w", entry.ModelID, err) + } + if ok { + return price, nil + } + } + if entry.PriceKnown { + return Price{ + InputMicroUSDPerMillion: entry.InputPriceMicroUSDPerMillion, + OutputMicroUSDPerMillion: entry.OutputPriceMicroUSDPerMillion, + Known: true, + }, nil + } + return Price{}, nil +} + +// bedrockPricesBucket is the cluster-replicated KV bucket holding per-model +// price overrides, keyed like bedrock-weights (base64url of the model ID, +// since model IDs contain ':'). +const bedrockPricesBucket = "bedrock-prices" + +// bedrockPricesHistory keeps one revision; a re-priced model overwrites in +// place. +const bedrockPricesHistory = 1 + +// PriceStore resolves per-model price overrides from the bedrock-prices +// JetStream KV bucket, mirroring WeightsStore's construction and lookup shape +// — both are per-model deployment data with an in-tree default plus a KV +// override (D4, D12). +type PriceStore struct { + js jetstream.JetStream + replicas int + + mu sync.Mutex + kv jetstream.KeyValue +} + +var _ PriceResolver = (*PriceStore)(nil) + +// NewPriceStore constructs a PriceStore over the cluster's JetStream client, +// replicated across replicas nodes. +func NewPriceStore(js jetstream.JetStream, replicas int) *PriceStore { + return &PriceStore{js: js, replicas: replicas} +} + +// bucket lazily opens (or creates) the cluster-replicated bedrock-prices KV +// bucket, caching the handle for subsequent calls, mirroring +// WeightsStore.bucket. +func (s *PriceStore) bucket(ctx context.Context) (jetstream.KeyValue, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.kv != nil { + return s.kv, nil + } + kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, s.js, bedrockPricesBucket, bedrockPricesHistory, s.replicas) + if err != nil { + return nil, err + } + s.kv = kv + return kv, nil +} + +// Resolve returns modelID's KV-overridden price, if one has been set. +func (s *PriceStore) Resolve(ctx context.Context, modelID string) (Price, bool, error) { + kv, err := s.bucket(ctx) + if err != nil { + return Price{}, false, err + } + entry, err := kv.Get(ctx, weightsKey(modelID)) + switch { + case err == nil: + var price Price + if err := json.Unmarshal(entry.Value(), &price); err != nil { + return Price{}, false, fmt.Errorf("decode price override for %s: %w", modelID, err) + } + return price, true, nil + case errors.Is(err, jetstream.ErrKeyNotFound): + return Price{}, false, nil + default: + return Price{}, false, fmt.Errorf("kv get price override for %s: %w", modelID, err) + } +} + +// PutPrice records price as modelID's overridden price, always Known +// (storing a KV override to mark a model unknown makes no sense — remove the +// override with DeletePrice instead to fall back to the in-tree default). +func (s *PriceStore) PutPrice(ctx context.Context, modelID string, price Price) error { + kv, err := s.bucket(ctx) + if err != nil { + return err + } + price.Known = true + data, err := json.Marshal(price) + if err != nil { + return fmt.Errorf("encode price override for %s: %w", modelID, err) + } + if _, err := kv.Put(ctx, weightsKey(modelID), data); err != nil { + return fmt.Errorf("kv put price override for %s: %w", modelID, err) + } + return nil +} + +// DeletePrice removes modelID's price override, falling resolution back to +// the catalog entry's in-tree default. Deleting an already-absent override is +// idempotent, not an error. +func (s *PriceStore) DeletePrice(ctx context.Context, modelID string) error { + kv, err := s.bucket(ctx) + if err != nil { + return err + } + if err := kv.Delete(ctx, weightsKey(modelID)); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + return fmt.Errorf("kv delete price override for %s: %w", modelID, err) + } + return nil +} diff --git a/spinifex/gateway/bedrock/prices_test.go b/spinifex/gateway/bedrock/prices_test.go new file mode 100644 index 000000000..7ac95a120 --- /dev/null +++ b/spinifex/gateway/bedrock/prices_test.go @@ -0,0 +1,123 @@ +package gateway_bedrock + +import ( + "context" + "testing" + + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + priceTestSelfHostModelID = "meta.llama3-2-1b-instruct-v1:0" + priceTestProviderModelID = "anthropic.claude-3-5-sonnet-20240620-v1:0" + priceTestUnknownModelID = "no-such-model" +) + +// stubPriceResolver returns a fixed (price, ok, err) for every Resolve call, +// standing in for a KV override without a real JetStream bucket. +type stubPriceResolver struct { + price Price + ok bool + err error +} + +func (s stubPriceResolver) Resolve(context.Context, string) (Price, bool, error) { + return s.price, s.ok, s.err +} + +// TestResolvePrice_SelfHost_AlwaysKnownZero asserts D12's hard rule: a +// self-hosted model's external cost is always a *known* zero — never +// consulted from the catalog's (irrelevant, zero-value) price fields — so it +// can never be confused with an unpriced model. +func TestResolvePrice_SelfHost_AlwaysKnownZero(t *testing.T) { + entry, ok := lookupCatalogEntry(priceTestSelfHostModelID) + require.True(t, ok) + + // A KV override present for this model must still be ignored: self-host + // pricing is a hard rule, not a resolvable-but-currently-zero default. + price, err := resolvePrice(context.Background(), stubPriceResolver{price: Price{InputMicroUSDPerMillion: 999, Known: true}, ok: true}, entry) + require.NoError(t, err) + assert.True(t, price.Known) + assert.Zero(t, price.InputMicroUSDPerMillion) + assert.Zero(t, price.OutputMicroUSDPerMillion) +} + +// TestResolvePrice_ProviderUsesInTreeDefault covers a provider entry with no +// KV override: it falls back to the catalog's in-tree default price. +func TestResolvePrice_ProviderUsesInTreeDefault(t *testing.T) { + entry, ok := lookupCatalogEntry(priceTestProviderModelID) + require.True(t, ok) + + price, err := resolvePrice(context.Background(), nil, entry) + require.NoError(t, err) + assert.True(t, price.Known) + assert.Equal(t, entry.InputPriceMicroUSDPerMillion, price.InputMicroUSDPerMillion) + assert.Equal(t, entry.OutputPriceMicroUSDPerMillion, price.OutputMicroUSDPerMillion) +} + +// TestResolvePrice_KVOverrideWinsOverInTreeDefault covers a provider entry +// with a KV override present: the override wins outright, not just fills a +// gap in the in-tree default. +func TestResolvePrice_KVOverrideWinsOverInTreeDefault(t *testing.T) { + entry, ok := lookupCatalogEntry(priceTestProviderModelID) + require.True(t, ok) + + override := Price{InputMicroUSDPerMillion: 1_000_000, OutputMicroUSDPerMillion: 2_000_000, Known: true} + price, err := resolvePrice(context.Background(), stubPriceResolver{price: override, ok: true}, entry) + require.NoError(t, err) + assert.Equal(t, override, price) + assert.NotEqual(t, entry.InputPriceMicroUSDPerMillion, price.InputMicroUSDPerMillion) +} + +// TestResolvePrice_NoDefaultNoOverride_Unknown is the D12 invariant under +// test: a model with neither a KV override nor an in-tree default resolves +// Known=false, never a silent zero. +func TestResolvePrice_NoDefaultNoOverride_Unknown(t *testing.T) { + entry := catalogEntry{ModelID: priceTestUnknownModelID, Provider: providerPrefix + "unpriced-vendor"} + + price, err := resolvePrice(context.Background(), nil, entry) + require.NoError(t, err) + assert.False(t, price.Known) +} + +// TestResolvePrice_ResolverErrorPropagates ensures a genuine KV read failure +// surfaces rather than silently falling through to the in-tree default or an +// unknown price — an internal fault is not a pricing verdict. +func TestResolvePrice_ResolverErrorPropagates(t *testing.T) { + entry, ok := lookupCatalogEntry(priceTestProviderModelID) + require.True(t, ok) + + _, err := resolvePrice(context.Background(), stubPriceResolver{err: assert.AnError}, entry) + require.Error(t, err) +} + +// TestPriceStore_PutGetDelete exercises the real JetStream KV path, mirroring +// TestWeightsStore_PutAndResolve_KV. +func TestPriceStore_PutGetDelete(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + store := NewPriceStore(testutil.NewJetStream(t, nc), 1) + ctx := context.Background() + + _, ok, err := store.Resolve(ctx, priceTestProviderModelID) + require.NoError(t, err) + assert.False(t, ok) + + require.NoError(t, store.PutPrice(ctx, priceTestProviderModelID, Price{InputMicroUSDPerMillion: 5_000_000, OutputMicroUSDPerMillion: 10_000_000})) + + got, ok, err := store.Resolve(ctx, priceTestProviderModelID) + require.NoError(t, err) + require.True(t, ok) + assert.True(t, got.Known, "PutPrice must always store a known price") + assert.EqualValues(t, 5_000_000, got.InputMicroUSDPerMillion) + assert.EqualValues(t, 10_000_000, got.OutputMicroUSDPerMillion) + + require.NoError(t, store.DeletePrice(ctx, priceTestProviderModelID)) + _, ok, err = store.Resolve(ctx, priceTestProviderModelID) + require.NoError(t, err) + assert.False(t, ok) + + // Deleting an already-absent override is idempotent, not an error. + require.NoError(t, store.DeletePrice(ctx, priceTestProviderModelID)) +} diff --git a/spinifex/gateway/bedrock/usage.go b/spinifex/gateway/bedrock/usage.go new file mode 100644 index 000000000..55f67b034 --- /dev/null +++ b/spinifex/gateway/bedrock/usage.go @@ -0,0 +1,440 @@ +package gateway_bedrock + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "strings" + "sync" + "time" + + "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/nats-io/nats.go/jetstream" +) + +const ( + // bedrockUsageBucket is the cluster-replicated KV bucket holding per + // account/model/period usage counters (D8, D12), keyed + // "{accountID}/{modelID-base64url}/{period}". + bedrockUsageBucket = "bedrock-usage" + bedrockUsageHistory = 1 + + // bedrockUsageDedupeBucket records every RequestID the usage consumer has + // already applied. At-least-once JetStream delivery means the same + // record can arrive more than once; kv.Create is atomic + // create-if-absent, so it doubles as the dedupe gate with no CAS loop of + // its own required. The TTL matches the invocation stream's own + // retention (invocationStreamMaxAge): once a record could no longer be + // redelivered, its dedupe marker has nothing left to guard against. + bedrockUsageDedupeBucket = "bedrock-usage-dedupe" + bedrockUsageDedupeHistory = 1 + + // usageCASRetries bounds Apply's retry on a revision conflict, mirroring + // handlers_quota's vcpuCASRetries. + usageCASRetries = 100 + + // UsageConsumerName is this package's durable pull consumer for usage and + // cost aggregation, distinct from InvocationDeliveryConsumer: both attach + // to the same OCHRE_INVOCATIONS stream (LimitsPolicy retention lets both + // see every message independent of the other's ack position) but own + // separate ack state, so redelivery on one never double-processes on the + // other. + UsageConsumerName = "ochre-usage-metering" +) + +// billingPeriod returns the monthly billing period key for t, e.g. "2026-08". +// Monthly is the only period this package understands; the tokens-per-month +// quota dimension reads the same key shape. +func billingPeriod(t time.Time) string { + return t.UTC().Format("2006-01") +} + +// usageKey returns the KV key for accountID's modelID usage in period. +// modelID is base64url-encoded via weightsKey (defined in weights.go, same +// package) since model IDs contain ':', which NATS rejects in a KV key. +func usageKey(accountID, modelID, period string) string { + return accountID + "/" + weightsKey(modelID) + "/" + period +} + +// UsageCounters is one account's accrued Bedrock usage for one model in one +// billing period. +type UsageCounters struct { + InputTokens int64 `json:"inputTokens"` + OutputTokens int64 `json:"outputTokens"` + RequestCount int64 `json:"requestCount"` + // CostMicroUSD is integer micro-USD (D12): floats drift across millions + // of increments, and drift in a billing counter is a customer dispute + // months later. Self-hosted contributions are always +0. + CostMicroUSD int64 `json:"costMicroUsd"` + // CostUnknown is sticky: once a contributing invocation had no + // resolvable price, CostMicroUSD under-counts that invocation forever, + // so the aggregate can never be called authoritative again for this + // period even if a later invocation's price does resolve. It never + // reverts to false once set, and must never be confused with a genuine + // zero-cost period (self-host, or a provider period with no usage). + CostUnknown bool `json:"costUnknown"` +} + +// UsageReader resolves accountID's total token usage (input+output, across +// every model) for the current billing period. Satisfied by UsageStore; the +// quota package depends on this shape without importing this package (see +// handlers_quota.BedrockUsageReader). +type UsageReader interface { + TokensThisPeriod(ctx context.Context, accountID string) (int64, error) +} + +// UsageStore persists per-account/model/period UsageCounters in the +// bedrock-usage JetStream KV bucket, and dedupes applied RequestIDs in a +// second bucket so at-least-once stream delivery never double-counts. +type UsageStore struct { + js jetstream.JetStream + replicas int + + mu sync.Mutex + kv jetstream.KeyValue + + dedupeMu sync.Mutex + dedupe jetstream.KeyValue +} + +var _ UsageReader = (*UsageStore)(nil) + +// NewUsageStore constructs a UsageStore over the cluster's JetStream client, +// replicated across replicas nodes. +func NewUsageStore(js jetstream.JetStream, replicas int) *UsageStore { + return &UsageStore{js: js, replicas: replicas} +} + +// bucket lazily opens (or creates) the cluster-replicated bedrock-usage KV +// bucket, caching the handle for subsequent calls. +func (s *UsageStore) bucket(ctx context.Context) (jetstream.KeyValue, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.kv != nil { + return s.kv, nil + } + kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, s.js, bedrockUsageBucket, bedrockUsageHistory, s.replicas) + if err != nil { + return nil, err + } + s.kv = kv + return kv, nil +} + +// dedupeBucket lazily opens (or creates) the cluster-replicated +// bedrock-usage-dedupe KV bucket. +func (s *UsageStore) dedupeBucket(ctx context.Context) (jetstream.KeyValue, error) { + s.dedupeMu.Lock() + defer s.dedupeMu.Unlock() + if s.dedupe != nil { + return s.dedupe, nil + } + kv, err := kvutil.GetOrCreateBucketWithTTL(ctx, s.js, bedrockUsageDedupeBucket, bedrockUsageDedupeHistory, invocationStreamMaxAge) + if err != nil { + return nil, err + } + s.dedupe = kv + return kv, nil +} + +// MarkProcessed atomically claims requestID for the usage consumer. first is +// true the one time this requestID is claimed; every subsequent call +// (redelivery) returns false with no error, so the caller can skip +// re-applying counters without treating the duplicate as a failure. +func (s *UsageStore) MarkProcessed(ctx context.Context, requestID string) (first bool, err error) { + kv, err := s.dedupeBucket(ctx) + if err != nil { + return false, err + } + if _, err := kv.Create(ctx, requestID, []byte("1")); err != nil { + if errors.Is(err, jetstream.ErrKeyExists) { + return false, nil + } + return false, fmt.Errorf("claim dedupe marker for %s: %w", requestID, err) + } + return true, nil +} + +// Get returns accountID's usage counters for modelID in period, or +// (zero, false, nil) if the account has accrued nothing yet. +func (s *UsageStore) Get(ctx context.Context, accountID, modelID, period string) (UsageCounters, bool, error) { + kv, err := s.bucket(ctx) + if err != nil { + return UsageCounters{}, false, err + } + entry, err := kv.Get(ctx, usageKey(accountID, modelID, period)) + switch { + case err == nil: + var counters UsageCounters + if uerr := json.Unmarshal(entry.Value(), &counters); uerr != nil { + return UsageCounters{}, false, fmt.Errorf("decode usage counters for %s/%s/%s: %w", accountID, modelID, period, uerr) + } + return counters, true, nil + case errors.Is(err, jetstream.ErrKeyNotFound): + return UsageCounters{}, false, nil + default: + return UsageCounters{}, false, fmt.Errorf("kv get usage counters for %s/%s/%s: %w", accountID, modelID, period, err) + } +} + +// microUSDCost returns tokens priced at pricePerMillion integer micro-USD +// per million tokens, truncated toward zero. Truncation is deliberate and +// consistent per record — it never accumulates in the customer's +// disfavour, and it keeps the arithmetic exact integer division with no +// float involved anywhere in the accrual path (D12). +func microUSDCost(tokens, pricePerMillion int64) int64 { + return tokens * pricePerMillion / 1_000_000 +} + +// Apply adds rec's token counts and (if price is known) accrued cost to its +// account/model/period counter, under CAS so concurrent consumers never lose +// an update. Callers MUST have already confirmed via MarkProcessed that this +// RequestID has not been applied before — Apply itself has no idempotency of +// its own and will double-count a record applied twice. +func (s *UsageStore) Apply(ctx context.Context, rec InvocationRecord, price Price) error { + kv, err := s.bucket(ctx) + if err != nil { + return err + } + key := usageKey(rec.AccountID, rec.ModelID, billingPeriod(rec.Timestamp)) + + for range usageCASRetries { + var current UsageCounters + var revision uint64 + entry, err := kv.Get(ctx, key) + switch { + case err == nil: + if uerr := json.Unmarshal(entry.Value(), ¤t); uerr != nil { + return fmt.Errorf("decode usage counters for %s: %w", key, uerr) + } + revision = entry.Revision() + case errors.Is(err, jetstream.ErrKeyNotFound): + // Zero value at revision 0: the CAS write below creates it. + default: + return fmt.Errorf("kv get usage counters for %s: %w", key, err) + } + + next := current + next.InputTokens += rec.InputTokens + next.OutputTokens += rec.OutputTokens + next.RequestCount++ + if price.Known { + next.CostMicroUSD += microUSDCost(rec.InputTokens, price.InputMicroUSDPerMillion) + microUSDCost(rec.OutputTokens, price.OutputMicroUSDPerMillion) + } else { + next.CostUnknown = true + } + + data, merr := json.Marshal(next) + if merr != nil { + return fmt.Errorf("encode usage counters for %s: %w", key, merr) + } + + if revision == 0 { + _, err = kv.Create(ctx, key, data) + } else { + _, err = kv.Update(ctx, key, data, revision) + } + if err == nil { + return nil + } + if !isUsageCASConflict(err) { + return fmt.Errorf("kv write usage counters for %s: %w", key, err) + } + } + return fmt.Errorf("usage counter CAS exhausted for %s after %d attempts", key, usageCASRetries) +} + +// TokensThisPeriod sums accountID's input+output tokens across every model +// for the current billing period, for the tokens-per-month quota dimension. +// A few seconds of staleness against the usage consumer is acceptable for a +// monthly cap (see handlers_quota.CheckBedrockTokens). +func (s *UsageStore) TokensThisPeriod(ctx context.Context, accountID string) (int64, error) { + kv, err := s.bucket(ctx) + if err != nil { + return 0, err + } + keys, err := kvutil.Keys(ctx, kv) + if err != nil { + if errors.Is(err, jetstream.ErrNoKeysFound) { + return 0, nil + } + return 0, fmt.Errorf("list usage keys for %s: %w", accountID, err) + } + + prefix := accountID + "/" + suffix := "/" + billingPeriod(time.Now()) + var total int64 + for _, key := range keys { + if !strings.HasPrefix(key, prefix) || !strings.HasSuffix(key, suffix) { + continue + } + entry, err := kv.Get(ctx, key) + if err != nil { + if errors.Is(err, jetstream.ErrKeyNotFound) { + continue + } + return 0, fmt.Errorf("kv get usage counters for %s: %w", key, err) + } + var counters UsageCounters + if uerr := json.Unmarshal(entry.Value(), &counters); uerr != nil { + return 0, fmt.Errorf("decode usage counters for %s: %w", key, uerr) + } + total += counters.InputTokens + counters.OutputTokens + } + return total, nil +} + +// isUsageCASConflict reports whether err is a lost optimistic-concurrency +// race on the usage counter: a Create on an existing key or an Update +// against a stale revision. Mirrors handlers_quota.isCASConflict — the two +// packages don't share an import for this one four-line check. +func isUsageCASConflict(err error) bool { + if errors.Is(err, jetstream.ErrKeyExists) { + return true + } + var apiErr *jetstream.APIError + if errors.As(err, &apiErr) { + return apiErr.ErrorCode == jetstream.JSErrCodeStreamWrongLastSequence + } + return false +} + +// EnsureUsageConsumer idempotently creates (or updates) this package's +// durable pull consumer for usage/cost metering on the invocation stream. +func EnsureUsageConsumer(ctx context.Context, js jetstream.JetStream) (jetstream.Consumer, error) { + consumer, err := js.CreateOrUpdateConsumer(ctx, InvocationStreamName, jetstream.ConsumerConfig{ + Durable: UsageConsumerName, + FilterSubject: InvocationStreamSubject, + AckPolicy: jetstream.AckExplicitPolicy, + MaxDeliver: 5, + AckWait: time.Minute, + }) + if err != nil { + return nil, fmt.Errorf("ensure usage metering consumer: %w", err) + } + return consumer, nil +} + +// UsageConsumer drains UsageConsumerName, dedupes on RequestID, and applies +// each record's token counts and priced cost to its account/model/period +// counter. +type UsageConsumer struct { + usage *UsageStore + prices PriceResolver +} + +// NewUsageConsumer constructs a UsageConsumer. prices may be nil, in which +// case every provider entry with no in-tree default price resolves unknown +// (only the in-tree defaults are ever consulted). +func NewUsageConsumer(usage *UsageStore, prices PriceResolver) *UsageConsumer { + return &UsageConsumer{usage: usage, prices: prices} +} + +// Run pulls records from consumer until ctx is cancelled, acking each after +// a successful apply (or a confirmed duplicate) and nak'ing on failure so +// JetStream redelivers it. +func (c *UsageConsumer) Run(ctx context.Context, consumer jetstream.Consumer) { + iter, err := consumer.Messages() + if err != nil { + slog.Error("bedrock: failed to start usage metering consumer", "err", err) + return + } + defer iter.Stop() + + // iter.Next() blocks indefinitely when no message is pending, so ctx + // cancellation alone would never unblock the loop below — Stop() is what + // actually wakes it, letting Next() return ErrMsgIteratorClosed. + stopped := make(chan struct{}) + defer close(stopped) + go func() { + select { + case <-ctx.Done(): + iter.Stop() + case <-stopped: + } + }() + + for { + msg, err := iter.Next() + if err != nil { + if errors.Is(err, jetstream.ErrMsgIteratorClosed) || ctx.Err() != nil { + return + } + slog.Error("bedrock: usage metering consumer fetch failed", "err", err) + continue + } + c.apply(ctx, msg) + } +} + +// apply decodes, dedupes, applies, and acks (or naks) one message. A decode +// failure acks and drops the message outright: it can never succeed on +// redelivery. +func (c *UsageConsumer) apply(ctx context.Context, msg jetstream.Msg) { + var rec InvocationRecord + if err := json.Unmarshal(msg.Data(), &rec); err != nil { + slog.Error("bedrock: usage metering consumer: undecodable record, dropping", "err", err) + if ackErr := msg.Ack(); ackErr != nil { + slog.Error("bedrock: usage metering consumer: ack failed", "err", ackErr) + } + return + } + + first, err := c.usage.MarkProcessed(ctx, rec.RequestID) + if err != nil { + slog.Error("bedrock: usage metering consumer: dedupe check failed", "request_id", rec.RequestID, "err", err) + if nakErr := msg.Nak(); nakErr != nil { + slog.Error("bedrock: usage metering consumer: nak failed", "err", nakErr) + } + return + } + if !first { + // Redelivery of a record this consumer already applied: ack without + // touching the counter, or a customer gets billed twice. + slog.Debug("bedrock: usage metering consumer: duplicate record, skipping", "request_id", rec.RequestID) + if ackErr := msg.Ack(); ackErr != nil { + slog.Error("bedrock: usage metering consumer: ack failed", "request_id", rec.RequestID, "err", ackErr) + } + return + } + + entry, ok := lookupCatalogEntry(rec.ModelID) + if !ok { + // A record for a model no longer in the catalog: still count tokens + // (they were genuinely consumed) but the cost can never be known. + slog.Warn("bedrock: usage metering consumer: record for unknown model, cost unresolved", "request_id", rec.RequestID, "model", rec.ModelID) + if err := c.usage.Apply(ctx, rec, Price{}); err != nil { + slog.Error("bedrock: usage metering consumer: apply failed", "request_id", rec.RequestID, "err", err) + if nakErr := msg.Nak(); nakErr != nil { + slog.Error("bedrock: usage metering consumer: nak failed", "err", nakErr) + } + return + } + if ackErr := msg.Ack(); ackErr != nil { + slog.Error("bedrock: usage metering consumer: ack failed", "request_id", rec.RequestID, "err", ackErr) + } + return + } + + price, err := resolvePrice(ctx, c.prices, entry) + if err != nil { + slog.Error("bedrock: usage metering consumer: price resolve failed", "request_id", rec.RequestID, "model", rec.ModelID, "err", err) + if nakErr := msg.Nak(); nakErr != nil { + slog.Error("bedrock: usage metering consumer: nak failed", "err", nakErr) + } + return + } + + if err := c.usage.Apply(ctx, rec, price); err != nil { + slog.Error("bedrock: usage metering consumer: apply failed", "request_id", rec.RequestID, "err", err) + if nakErr := msg.Nak(); nakErr != nil { + slog.Error("bedrock: usage metering consumer: nak failed", "err", nakErr) + } + return + } + if err := msg.Ack(); err != nil { + slog.Error("bedrock: usage metering consumer: ack failed", "request_id", rec.RequestID, "err", err) + } +} diff --git a/spinifex/gateway/bedrock/usage_test.go b/spinifex/gateway/bedrock/usage_test.go new file mode 100644 index 000000000..44efe4001 --- /dev/null +++ b/spinifex/gateway/bedrock/usage_test.go @@ -0,0 +1,290 @@ +package gateway_bedrock + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + usageTestSelfHostModelID = "meta.llama3-2-1b-instruct-v1:0" + usageTestProviderModelID = "anthropic.claude-3-5-sonnet-20240620-v1:0" + usageTestAccountA = "acct-a" + usageTestAccountB = "acct-b" +) + +func newTestUsageStore(t *testing.T) *UsageStore { + t.Helper() + _, nc, _ := testutil.StartTestJetStream(t) + return NewUsageStore(testutil.NewJetStream(t, nc), 1) +} + +// TestMicroUSDCost_Arithmetic covers D12's integer accrual, including +// truncation for token counts too small to register a whole micro-USD. +func TestMicroUSDCost_Arithmetic(t *testing.T) { + tests := []struct { + name string + tokens int64 + pricePerMillion int64 + want int64 + }{ + {"one million tokens at $3/MTok", 1_000_000, 3_000_000, 3_000_000}, + {"1.5 million tokens", 1_500_000, 3_000_000, 4_500_000}, + {"zero tokens", 0, 3_000_000, 0}, + {"sub-unit truncates toward zero", 1, 3_000_000, 3}, + {"one token at $1/MTok truncates to zero", 1, 1, 0}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, microUSDCost(tt.tokens, tt.pricePerMillion)) + }) + } +} + +// TestUsageStore_Apply_AccumulatesAcrossCalls covers the core counter shape: +// two records for the same account/model/period sum tokens, request count, +// and cost rather than overwriting. +func TestUsageStore_Apply_AccumulatesAcrossCalls(t *testing.T) { + store := newTestUsageStore(t) + ctx := context.Background() + now := time.Now() + price := Price{InputMicroUSDPerMillion: 3_000_000, OutputMicroUSDPerMillion: 15_000_000, Known: true} + + rec1 := InvocationRecord{AccountID: usageTestAccountA, ModelID: usageTestProviderModelID, Timestamp: now, InputTokens: 1_000_000, OutputTokens: 500_000} + rec2 := InvocationRecord{AccountID: usageTestAccountA, ModelID: usageTestProviderModelID, Timestamp: now, InputTokens: 2_000_000, OutputTokens: 1_000_000} + + require.NoError(t, store.Apply(ctx, rec1, price)) + require.NoError(t, store.Apply(ctx, rec2, price)) + + counters, ok, err := store.Get(ctx, usageTestAccountA, usageTestProviderModelID, billingPeriod(now)) + require.NoError(t, err) + require.True(t, ok) + assert.EqualValues(t, 3_000_000, counters.InputTokens) + assert.EqualValues(t, 1_500_000, counters.OutputTokens) + assert.EqualValues(t, 2, counters.RequestCount) + // (1M*3 + 0.5M*15) + (2M*3 + 1M*15) = (3M + 7.5M) + (6M + 15M) = 31.5M + assert.EqualValues(t, 31_500_000, counters.CostMicroUSD) + assert.False(t, counters.CostUnknown) +} + +// TestUsageStore_Apply_SelfHostZeroCostTokensRecorded is the D12 self-host +// rule under test: tokens are recorded exactly like a paid model, but cost +// accrues a *known* zero, never a flag-less absence. +func TestUsageStore_Apply_SelfHostZeroCostTokensRecorded(t *testing.T) { + store := newTestUsageStore(t) + ctx := context.Background() + now := time.Now() + + entry, ok := lookupCatalogEntry(usageTestSelfHostModelID) + require.True(t, ok) + price, err := resolvePrice(ctx, nil, entry) + require.NoError(t, err) + + rec := InvocationRecord{AccountID: usageTestAccountA, ModelID: usageTestSelfHostModelID, Timestamp: now, InputTokens: 42, OutputTokens: 17} + require.NoError(t, store.Apply(ctx, rec, price)) + + counters, ok, err := store.Get(ctx, usageTestAccountA, usageTestSelfHostModelID, billingPeriod(now)) + require.NoError(t, err) + require.True(t, ok) + assert.EqualValues(t, 42, counters.InputTokens) + assert.EqualValues(t, 17, counters.OutputTokens) + assert.EqualValues(t, 0, counters.CostMicroUSD) + assert.False(t, counters.CostUnknown, "self-host must accrue a KNOWN zero, not unknown") +} + +// TestUsageStore_Apply_UnknownPriceFlagsCostUnknown_NotZero is the other half +// of the same invariant: an unresolvable price must never look like a $0 +// period. CostMicroUSD stays whatever it already was; CostUnknown is the only +// signal a reader may trust. +func TestUsageStore_Apply_UnknownPriceFlagsCostUnknown_NotZero(t *testing.T) { + store := newTestUsageStore(t) + ctx := context.Background() + now := time.Now() + + rec := InvocationRecord{AccountID: usageTestAccountA, ModelID: usageTestProviderModelID, Timestamp: now, InputTokens: 100, OutputTokens: 50} + require.NoError(t, store.Apply(ctx, rec, Price{})) // Known: false + + counters, ok, err := store.Get(ctx, usageTestAccountA, usageTestProviderModelID, billingPeriod(now)) + require.NoError(t, err) + require.True(t, ok) + assert.EqualValues(t, 100, counters.InputTokens) + assert.EqualValues(t, 50, counters.OutputTokens) + assert.EqualValues(t, 0, counters.CostMicroUSD) + assert.True(t, counters.CostUnknown) + + // A later record with a KNOWN price must not clear the flag: the period's + // total is permanently missing the earlier unpriced contribution. + require.NoError(t, store.Apply(ctx, rec, Price{InputMicroUSDPerMillion: 3_000_000, OutputMicroUSDPerMillion: 15_000_000, Known: true})) + counters, ok, err = store.Get(ctx, usageTestAccountA, usageTestProviderModelID, billingPeriod(now)) + require.NoError(t, err) + require.True(t, ok) + assert.True(t, counters.CostUnknown, "CostUnknown must stay sticky once set") + assert.Positive(t, counters.CostMicroUSD, "the priced record's cost is still added") +} + +// TestUsageStore_MarkProcessed_DedupesRedelivery is the unit-level half of +// the mandatory dedupe: the same RequestID may only be claimed once. +func TestUsageStore_MarkProcessed_DedupesRedelivery(t *testing.T) { + store := newTestUsageStore(t) + ctx := context.Background() + + first, err := store.MarkProcessed(ctx, "req-1") + require.NoError(t, err) + assert.True(t, first) + + second, err := store.MarkProcessed(ctx, "req-1") + require.NoError(t, err) + assert.False(t, second, "a redelivered RequestID must not be claimed twice") + + // A distinct RequestID is unaffected. + third, err := store.MarkProcessed(ctx, "req-2") + require.NoError(t, err) + assert.True(t, third) +} + +// TestUsageStore_TokensThisPeriod_SumsAcrossModels covers the reader the +// tokens-per-month quota dimension consults: total is input+output summed +// across every model for the account, in the current period only. +func TestUsageStore_TokensThisPeriod_SumsAcrossModels(t *testing.T) { + store := newTestUsageStore(t) + ctx := context.Background() + now := time.Now() + + require.NoError(t, store.Apply(ctx, InvocationRecord{AccountID: usageTestAccountA, ModelID: usageTestSelfHostModelID, Timestamp: now, InputTokens: 100, OutputTokens: 50}, Price{Known: true})) + require.NoError(t, store.Apply(ctx, InvocationRecord{AccountID: usageTestAccountA, ModelID: usageTestProviderModelID, Timestamp: now, InputTokens: 200, OutputTokens: 25}, Price{Known: true})) + // A different account's usage must not bleed into accountA's total. + require.NoError(t, store.Apply(ctx, InvocationRecord{AccountID: usageTestAccountB, ModelID: usageTestSelfHostModelID, Timestamp: now, InputTokens: 9_999, OutputTokens: 9_999}, Price{Known: true})) + + total, err := store.TokensThisPeriod(ctx, usageTestAccountA) + require.NoError(t, err) + assert.EqualValues(t, 100+50+200+25, total) +} + +// TestUsageStore_TokensThisPeriod_NoUsageIsZero covers the empty-bucket path +// (jetstream.ErrNoKeysFound) rather than erroring. +func TestUsageStore_TokensThisPeriod_NoUsageIsZero(t *testing.T) { + store := newTestUsageStore(t) + total, err := store.TokensThisPeriod(context.Background(), usageTestAccountA) + require.NoError(t, err) + assert.Zero(t, total) +} + +// TestUsageConsumer_Run_RedeliveredRecordDoesNotDoubleCount is the +// bead-mandated assertion: an at-least-once redelivery of the same +// InvocationRecord — modelled here as the same JSON payload landing on the +// stream twice, the shape a StreamRecorder retry after an ambiguous publish +// outcome takes on the wire — must be applied exactly once. +func TestUsageConsumer_Run_RedeliveredRecordDoesNotDoubleCount(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + _, err := EnsureInvocationStream(ctx, js, 1) + require.NoError(t, err) + consumer, err := EnsureUsageConsumer(ctx, js) + require.NoError(t, err) + + now := time.Now() + rec := InvocationRecord{ + RequestID: "req-redelivered", AccountID: usageTestAccountA, ModelID: usageTestSelfHostModelID, + Operation: OperationConverse, Timestamp: now, InputTokens: 10, OutputTokens: 5, + } + payload, err := json.Marshal(rec) + require.NoError(t, err) + + // Two distinct stream messages carrying the same RequestID: JetStream's + // own publish-time msgID dedup only guards a single Publish call, not + // this scenario, which is exactly why a consumer-side dedupe is + // mandatory. + _, err = js.Publish(ctx, InvocationStreamSubject, payload) + require.NoError(t, err) + _, err = js.Publish(ctx, InvocationStreamSubject, payload) + require.NoError(t, err) + + usage := NewUsageStore(js, 1) + uc := NewUsageConsumer(usage, nil) + done := make(chan struct{}) + go func() { + uc.Run(ctx, consumer) + close(done) + }() + + period := billingPeriod(now) + require.Eventually(t, func() bool { + counters, ok, err := usage.Get(context.Background(), usageTestAccountA, usageTestSelfHostModelID, period) + return err == nil && ok && counters.RequestCount >= 1 + }, 2*time.Second, 10*time.Millisecond) + + // Give the (already-processed) second message a moment to be consumed + // and confirmed as a duplicate before asserting the final count. + time.Sleep(200 * time.Millisecond) + cancel() + <-done + + counters, ok, err := usage.Get(context.Background(), usageTestAccountA, usageTestSelfHostModelID, period) + require.NoError(t, err) + require.True(t, ok) + assert.EqualValues(t, 1, counters.RequestCount, "redelivery must not double-count the request") + assert.EqualValues(t, 10, counters.InputTokens, "redelivery must not double-count tokens") + assert.EqualValues(t, 5, counters.OutputTokens) +} + +// TestUsageConsumer_Run_UnknownModelStillRecordsTokensCostUnknown covers a +// record for a model the catalog no longer knows (e.g. retired since the +// call was made): tokens were genuinely consumed and must still be counted, +// but cost can never be resolved for it. +func TestUsageConsumer_Run_UnknownModelStillRecordsTokensCostUnknown(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + _, err := EnsureInvocationStream(ctx, js, 1) + require.NoError(t, err) + consumer, err := EnsureUsageConsumer(ctx, js) + require.NoError(t, err) + + now := time.Now() + rec := InvocationRecord{ + RequestID: "req-retired-model", AccountID: usageTestAccountA, ModelID: "retired.model-v0:0", + Operation: OperationConverse, Timestamp: now, InputTokens: 8, OutputTokens: 4, + } + payload, err := json.Marshal(rec) + require.NoError(t, err) + _, err = js.Publish(ctx, InvocationStreamSubject, payload) + require.NoError(t, err) + + usage := NewUsageStore(js, 1) + uc := NewUsageConsumer(usage, nil) + done := make(chan struct{}) + go func() { + uc.Run(ctx, consumer) + close(done) + }() + + period := billingPeriod(now) + require.Eventually(t, func() bool { + _, ok, err := usage.Get(context.Background(), usageTestAccountA, "retired.model-v0:0", period) + return err == nil && ok + }, 2*time.Second, 10*time.Millisecond) + cancel() + <-done + + counters, ok, err := usage.Get(context.Background(), usageTestAccountA, "retired.model-v0:0", period) + require.NoError(t, err) + require.True(t, ok) + assert.EqualValues(t, 8, counters.InputTokens) + assert.EqualValues(t, 4, counters.OutputTokens) + assert.True(t, counters.CostUnknown) +} + +// TestBillingPeriod_MonthlyFormat pins the period key shape, since +// TokensThisPeriod's prefix/suffix matching depends on it exactly. +func TestBillingPeriod_MonthlyFormat(t *testing.T) { + ts := time.Date(2026, time.August, 3, 12, 0, 0, 0, time.UTC) + assert.Equal(t, "2026-08", billingPeriod(ts)) +} diff --git a/spinifex/gateway/bedrockruntime.go b/spinifex/gateway/bedrockruntime.go index da934a1c9..4f73536a6 100644 --- a/spinifex/gateway/bedrockruntime.go +++ b/spinifex/gateway/bedrockruntime.go @@ -103,6 +103,17 @@ func (gw *GatewayConfig) BedrockRuntime_Request(w http.ResponseWriter, r *http.R return errors.New(awserrors.ErrorServerInternal) } + // Ochre metering (mulga-vmfng.7.6). RPM is a local, in-memory check — + // cheap enough to run before the body is even read. The token cap reads + // the stream-fed usage counter, so it goes second; both are no-ops + // (nil-safe on gw.Quota) unless their dimension is explicitly enabled. + if err := gw.Quota.CheckBedrockRPM(accountID); err != nil { + return err + } + if err := gw.Quota.CheckBedrockTokens(r.Context(), accountID); err != nil { + return err + } + body, err := io.ReadAll(r.Body) if err != nil { slog.ErrorContext(r.Context(), "BedrockRuntime_Request: failed to read body", "err", err) diff --git a/spinifex/handlers/quota/bedrock.go b/spinifex/handlers/quota/bedrock.go new file mode 100644 index 000000000..29c7a703b --- /dev/null +++ b/spinifex/handlers/quota/bedrock.go @@ -0,0 +1,257 @@ +package handlers_quota + +import ( + "context" + "encoding/json" + "errors" + "log/slog" + "sync" + "time" + + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/nats-io/nats.go/jetstream" +) + +// BedrockUsageReader resolves accountID's current-period Bedrock token usage +// (input+output, across every model), the counter gateway_bedrock's usage +// consumer maintains from the invocation stream (mulga-vmfng.7.6). The +// interface is defined here rather than importing gateway_bedrock so the two +// packages stay decoupled; gateway_bedrock.UsageStore satisfies it +// structurally, and awsgw.go wires the two together at startup. +type BedrockUsageReader interface { + TokensThisPeriod(ctx context.Context, accountID string) (int64, error) +} + +// SetBedrockUsage installs the reader CheckBedrockTokens consults. Called +// once at gateway startup; a nil (never-installed) reader leaves the +// dimension exempt even when TokensPerMonthEnabled — a wiring gap must not +// block every Bedrock call. +func (s *Service) SetBedrockUsage(r BedrockUsageReader) { + s.bedrockUsage = r +} + +// CheckBedrockTokens rejects with ServiceQuotaExceededException when +// accountID's current-period token usage has already reached the configured +// monthly cap. Unlike CheckVCPU's check-before-grant shape, a call's token +// cost is only known once it completes, so this is a check-against-history +// gate: the invocation that crosses the cap is itself allowed (its cost +// isn't known until it finishes), and the next one is rejected once the +// stream-fed counter catches up. A few seconds of lag behind real usage is +// an accepted property of a monthly cap (D8). +func (s *Service) CheckBedrockTokens(ctx context.Context, accountID string) error { + if s == nil || !s.limits.TokensPerMonthEnabled || accountID == utils.GlobalAccountID { + return nil + } + if s.bedrockUsage == nil { + return nil + } + tokens, err := s.bedrockUsage.TokensThisPeriod(ctx, accountID) + if err != nil { + return err + } + if tokens >= s.limits.TokensPerMonth { + return errors.New(awserrors.ErrorServiceQuotaExceededException) + } + return nil +} + +// CheckBedrockRPM rejects with ThrottlingException when accountID has +// exhausted its local per-gateway request-rate bucket. Enforcement is local +// and immediate — it never reads the invocation stream or a shared KV +// counter — which is deliberately generous under multiple gateways (N +// gateways x bucket capacity each): the correct direction to err for a rate +// limit is to occasionally allow a little more, never to add network latency +// to every invocation for a slightly tighter cap. +func (s *Service) CheckBedrockRPM(accountID string) error { + if s == nil || !s.limits.RequestsPerMinuteEnabled || accountID == utils.GlobalAccountID { + return nil + } + if !s.rpm.allow(accountID, s.limits.RequestsPerMinute) { + return errors.New(awserrors.ErrorThrottlingException) + } + return nil +} + +// tokenBucket is a classic token-bucket rate limiter: capacity tokens +// refilled continuously at capacity/60 per second, consumed one at a time. +// Refill is lazy — computed from elapsed wall-clock time on each Allow call — +// so no background goroutine is required per account. +type tokenBucket struct { + mu sync.Mutex + tokens float64 + capacity float64 + refillRate float64 // tokens per second + last time.Time + now func() time.Time +} + +// newTokenBucket constructs a full bucket (a fresh account starts able to +// burst its whole per-minute allowance, not throttled from zero) at +// capacityPerMinute, using now for its clock — overridable in tests, nil +// defaults to time.Now. +func newTokenBucket(capacityPerMinute int, now func() time.Time) *tokenBucket { + if now == nil { + now = time.Now + } + capacity := float64(max(capacityPerMinute, 0)) + return &tokenBucket{ + tokens: capacity, + capacity: capacity, + refillRate: capacity / 60, + last: now(), + now: now, + } +} + +// Allow refills for elapsed time since the last call, then attempts to +// consume one token. Returns false (throttle) when the bucket is empty. +func (b *tokenBucket) Allow() bool { + b.mu.Lock() + defer b.mu.Unlock() + b.refillLocked() + if b.tokens < 1 { + return false + } + b.tokens-- + return true +} + +// remaining reports the bucket's current occupancy (refilled up to now) +// without consuming a token, for the periodic KV sync's observability +// snapshot. +func (b *tokenBucket) remaining() int { + b.mu.Lock() + defer b.mu.Unlock() + b.refillLocked() + return int(b.tokens) +} + +// refillLocked adds tokens for the time elapsed since last, capped at +// capacity. Caller must hold mu. +func (b *tokenBucket) refillLocked() { + current := b.now() + elapsed := current.Sub(b.last).Seconds() + if elapsed <= 0 { + return + } + b.tokens = min(b.capacity, b.tokens+elapsed*b.refillRate) + b.last = current +} + +// rpmLimiter holds one tokenBucket per account, created lazily on first use +// at the configured per-minute capacity. Buckets are never evicted: an idle +// account costs one small map entry, never unbounded growth, and a restart +// resets every bucket to full — acceptable since RPM enforcement is already +// documented as generous, not exact. +type rpmLimiter struct { + mu sync.Mutex + buckets map[string]*tokenBucket + // now overrides tokenBucket's clock for every bucket this limiter + // creates; nil (the production default) leaves each bucket on time.Now. + now func() time.Time +} + +func newRPMLimiter() *rpmLimiter { + return &rpmLimiter{buckets: make(map[string]*tokenBucket)} +} + +// allow returns whether accountID may proceed under capacityPerMinute, +// creating its bucket on first use. +func (l *rpmLimiter) allow(accountID string, capacityPerMinute int) bool { + return l.bucketFor(accountID, capacityPerMinute).Allow() +} + +// bucketFor returns accountID's bucket, creating it at capacityPerMinute on +// first use. Capacity is fixed at creation: RPM limits are process-lifetime +// config (loaded once at gateway startup), not hot-reloaded. +func (l *rpmLimiter) bucketFor(accountID string, capacityPerMinute int) *tokenBucket { + l.mu.Lock() + defer l.mu.Unlock() + b, ok := l.buckets[accountID] + if !ok { + b = newTokenBucket(capacityPerMinute, l.now) + l.buckets[accountID] = b + } + return b +} + +// snapshot returns every known account's current bucket occupancy, for the +// periodic KV sync below. +func (l *rpmLimiter) snapshot() map[string]int { + l.mu.Lock() + defer l.mu.Unlock() + out := make(map[string]int, len(l.buckets)) + for accountID, b := range l.buckets { + out[accountID] = b.remaining() + } + return out +} + +// KVBucketBedrockRPMUsage is the gateway-owned KV bucket RunBedrockRPMSync +// writes each account's current bucket occupancy to, keyed by accountID. It +// exists for cross-node visibility and dashboards only: CheckBedrockRPM never +// reads it back. Enforcement must stay local and immediate (D8) — reading a +// remote value here would reintroduce exactly the round-trip latency that +// ruled out stream-fed RPM enforcement in the first place. Concurrent +// gateways overwrite each other's snapshot for the same account with no +// coordination, which is fine: this bucket is observational, not +// correctness-bearing. +const KVBucketBedrockRPMUsage = "bedrock-rpm-usage" + +// bedrockRPMSyncInterval is the sync cadence, matching the vCPU reconcile's +// ReconcileInterval so the two background loops share one cadence to reason +// about. +const bedrockRPMSyncInterval = ReconcileInterval + +// bedrockRPMSnapshot is one account's RPM bucket occupancy at sync time. +type bedrockRPMSnapshot struct { + Remaining int `json:"remaining"` + Capacity int `json:"capacity"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// RunBedrockRPMSync periodically snapshots every account this gateway has +// rate-limited into kv, until ctx is cancelled. A no-op when the RPM +// dimension is disabled, so a default-off deployment starts no ticker. +func (s *Service) RunBedrockRPMSync(ctx context.Context, js jetstream.JetStream, replicas int) { + if s == nil || !s.limits.RequestsPerMinuteEnabled { + return + } + kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, js, KVBucketBedrockRPMUsage, 1, replicas) + if err != nil { + slog.Warn("bedrock rpm sync: open usage bucket failed, sync disabled", "err", err) + return + } + ticker := time.NewTicker(bedrockRPMSyncInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + s.syncBedrockRPMOnce(ctx, kv) + } + } +} + +// syncBedrockRPMOnce writes one snapshot pass of every known account's +// bucket occupancy to kv. A per-account write failure is logged and the pass +// continues, since this data is observational only. +func (s *Service) syncBedrockRPMOnce(ctx context.Context, kv jetstream.KeyValue) { + for accountID, remaining := range s.rpm.snapshot() { + data, err := json.Marshal(bedrockRPMSnapshot{ + Remaining: remaining, + Capacity: s.limits.RequestsPerMinute, + UpdatedAt: time.Now(), + }) + if err != nil { + slog.Warn("bedrock rpm sync: encode snapshot failed", "account", accountID, "err", err) + continue + } + if _, err := kv.Put(ctx, accountID, data); err != nil { + slog.Warn("bedrock rpm sync: write failed", "account", accountID, "err", err) + } + } +} diff --git a/spinifex/handlers/quota/bedrock_test.go b/spinifex/handlers/quota/bedrock_test.go new file mode 100644 index 000000000..cf96a1844 --- /dev/null +++ b/spinifex/handlers/quota/bedrock_test.go @@ -0,0 +1,232 @@ +package handlers_quota + +import ( + "context" + "encoding/json" + "errors" + "testing" + "time" + + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/nats-io/nats.go/jetstream" + toml "github.com/pelletier/go-toml/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// stubBedrockUsageReader returns a fixed (tokens, err) for every account. +type stubBedrockUsageReader struct { + tokens int64 + err error +} + +func (s stubBedrockUsageReader) TokensThisPeriod(context.Context, string) (int64, error) { + return s.tokens, s.err +} + +// TestCheckBedrockTokens_Boundaries covers the hard cap at, just under, and +// just over the configured monthly limit. +func TestCheckBedrockTokens_Boundaries(t *testing.T) { + tests := []struct { + name string + tokens int64 + wantExceed bool + }{ + {"just under cap passes", 999, false}, + {"at cap rejects", 1000, true}, + {"over cap rejects", 1001, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := New(Limits{TokensPerMonthEnabled: true, TokensPerMonth: 1000}, nil) + s.SetBedrockUsage(stubBedrockUsageReader{tokens: tt.tokens}) + err := s.CheckBedrockTokens(context.Background(), testAccount) + if !tt.wantExceed { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Equal(t, awserrors.ErrorServiceQuotaExceededException, err.Error()) + }) + } +} + +func TestCheckBedrockTokens_DisabledIsNoop(t *testing.T) { + s := New(Limits{}, nil) + s.SetBedrockUsage(stubBedrockUsageReader{tokens: 1_000_000_000}) + assert.NoError(t, s.CheckBedrockTokens(context.Background(), testAccount)) +} + +func TestCheckBedrockTokens_ExemptSystemAccount(t *testing.T) { + s := New(Limits{TokensPerMonthEnabled: true, TokensPerMonth: 1}, nil) + s.SetBedrockUsage(stubBedrockUsageReader{tokens: 1_000_000}) + assert.NoError(t, s.CheckBedrockTokens(context.Background(), utils.GlobalAccountID)) +} + +// TestCheckBedrockTokens_NoReaderIsNoop covers a wiring gap: the dimension is +// enabled but SetBedrockUsage was never called. Failing open here (rather +// than blocking every Bedrock call) is a deliberate choice, distinct from a +// live read failure, which fails closed below. +func TestCheckBedrockTokens_NoReaderIsNoop(t *testing.T) { + s := New(Limits{TokensPerMonthEnabled: true, TokensPerMonth: 1}, nil) + assert.NoError(t, s.CheckBedrockTokens(context.Background(), testAccount)) +} + +// TestCheckBedrockTokens_ReaderErrorPropagates covers a live read failure +// (e.g. the usage KV bucket is unreachable): this fails closed, surfacing +// the error rather than silently allowing the call through. +func TestCheckBedrockTokens_ReaderErrorPropagates(t *testing.T) { + s := New(Limits{TokensPerMonthEnabled: true, TokensPerMonth: 1000}, nil) + s.SetBedrockUsage(stubBedrockUsageReader{err: errors.New("boom")}) + err := s.CheckBedrockTokens(context.Background(), testAccount) + require.Error(t, err) + assert.Equal(t, "boom", err.Error()) +} + +// TestCheckBedrockTokens_NilServiceIsNoop mirrors CheckVCPU's nil-safety: +// route handlers call gw.Quota.CheckBedrockTokens unconditionally, and +// gw.Quota is nil in test harnesses/unconfigured gateways. +func TestCheckBedrockTokens_NilServiceIsNoop(t *testing.T) { + var s *Service + assert.NoError(t, s.CheckBedrockTokens(context.Background(), testAccount)) +} + +func TestCheckBedrockRPM_DisabledIsNoop(t *testing.T) { + s := New(Limits{}, nil) + for range 1000 { + assert.NoError(t, s.CheckBedrockRPM(testAccount)) + } +} + +func TestCheckBedrockRPM_ExemptSystemAccount(t *testing.T) { + s := New(Limits{RequestsPerMinuteEnabled: true, RequestsPerMinute: 1}, nil) + for range 10 { + assert.NoError(t, s.CheckBedrockRPM(utils.GlobalAccountID)) + } +} + +func TestCheckBedrockRPM_NilServiceIsNoop(t *testing.T) { + var s *Service + assert.NoError(t, s.CheckBedrockRPM(testAccount)) +} + +// TestCheckBedrockRPM_ThrottlesAfterCapacityExhausted covers immediate +// enforcement: the (capacity+1)th request within the same instant is +// throttled, and a different account is unaffected. +func TestCheckBedrockRPM_ThrottlesAfterCapacityExhausted(t *testing.T) { + s := New(Limits{RequestsPerMinuteEnabled: true, RequestsPerMinute: 3}, nil) + + for i := range 3 { + assert.NoErrorf(t, s.CheckBedrockRPM(testAccount), "request %d within capacity", i) + } + err := s.CheckBedrockRPM(testAccount) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorThrottlingException, err.Error()) + + // A different account has its own, unexhausted bucket. + assert.NoError(t, s.CheckBedrockRPM("other-account")) +} + +// TestTokenBucket_RefillsOverTime drives the bucket with a fake clock so +// refill behaviour is deterministic rather than sleep-based. +func TestTokenBucket_RefillsOverTime(t *testing.T) { + now := time.Now() + clock := func() time.Time { return now } + b := newTokenBucket(60, clock) // 60/min == 1/sec + + // Drain the bucket. + for range 60 { + require.True(t, b.Allow()) + } + assert.False(t, b.Allow(), "bucket must be empty immediately after draining") + + // Advance 5 seconds: 5 tokens should have refilled. + now = now.Add(5 * time.Second) + for i := range 5 { + assert.Truef(t, b.Allow(), "refilled token %d", i) + } + assert.False(t, b.Allow(), "no more than the refilled amount is available") + + // Advance well past capacity: refill caps at capacity, it never + // overflows from an idle account banking unlimited tokens. + now = now.Add(10 * time.Minute) + assert.Equal(t, 60, b.remaining()) +} + +// TestRPMLimiter_PerAccountIsolation covers that one account's bucket never +// interferes with another's. +func TestRPMLimiter_PerAccountIsolation(t *testing.T) { + l := newRPMLimiter() + assert.True(t, l.allow("a", 1)) + assert.False(t, l.allow("a", 1)) + assert.True(t, l.allow("b", 1)) +} + +// TestService_SyncBedrockRPMOnce_WritesSnapshotToKV covers the periodic KV +// sync: after consuming from an account's bucket, a sync pass writes its +// current occupancy to KV for cross-node observability. +func TestService_SyncBedrockRPMOnce_WritesSnapshotToKV(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + js := testutil.NewJetStream(t, nc) + kv, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{Bucket: KVBucketBedrockRPMUsage, History: 1}) + require.NoError(t, err) + + s := New(Limits{RequestsPerMinuteEnabled: true, RequestsPerMinute: 10}, nil) + require.NoError(t, s.CheckBedrockRPM(testAccount)) + require.NoError(t, s.CheckBedrockRPM(testAccount)) + + s.syncBedrockRPMOnce(t.Context(), kv) + + entry, err := kv.Get(t.Context(), testAccount) + require.NoError(t, err) + var snapshot bedrockRPMSnapshot + require.NoError(t, json.Unmarshal(entry.Value(), &snapshot)) + assert.Equal(t, 8, snapshot.Remaining) + assert.Equal(t, 10, snapshot.Capacity) +} + +// TestRunBedrockRPMSync_DisabledNoop asserts the sync loop never touches its +// (here nil) JetStream client when the dimension is disabled — a real +// deployment leaves js valid, but the short-circuit must come first. +func TestRunBedrockRPMSync_DisabledNoop(t *testing.T) { + s := New(Limits{}, nil) + done := make(chan struct{}) + go func() { + s.RunBedrockRPMSync(context.Background(), nil, 1) + close(done) + }() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("RunBedrockRPMSync did not return promptly when disabled") + } +} + +// TestLimitsTOMLRoundTrip_BedrockDimensions extends the base round-trip test +// to the two new Ochre dimensions, including confirming they default to +// disabled when absent from the TOML source. +func TestLimitsTOMLRoundTrip_BedrockDimensions(t *testing.T) { + const src = ` +[quota] +enabled = true +vcpus = 8 +tokens_per_month_enabled = true +tokens_per_month = 1000000 +requests_per_minute_enabled = true +requests_per_minute = 60 +` + want := Limits{ + Enabled: true, + VCPUs: 8, + TokensPerMonthEnabled: true, + TokensPerMonth: 1_000_000, + RequestsPerMinuteEnabled: true, + RequestsPerMinute: 60, + } + + var parsed quotaTOML + require.NoError(t, toml.Unmarshal([]byte(src), &parsed)) + assert.Equal(t, want, parsed.Quota) +} diff --git a/spinifex/handlers/quota/quota.go b/spinifex/handlers/quota/quota.go index 0fa16c123..928401bf2 100644 --- a/spinifex/handlers/quota/quota.go +++ b/spinifex/handlers/quota/quota.go @@ -27,6 +27,14 @@ const KVBucketQuotaReconcile = "spinifex-quota-reconcile" // Limits mirrors the [quota] block in awsgw.toml. The zero value (Enabled // false) is a valid no-op, so gateways without a [quota] block are unaffected. +// +// TokensPerMonthEnabled and RequestsPerMinuteEnabled are deliberately +// separate from Enabled above (mulga-vmfng.7.6): a deployment already +// enforcing standing-infra quotas (Enabled=true for VCPUs/VPCs/...) must opt +// in to Bedrock token/RPM enforcement independently, since a bare Enabled +// would otherwise silently cap Bedrock usage at zero the moment any other +// dimension turns on. Both default to disabled, matching the block's +// existing opt-in convention. type Limits struct { Enabled bool `toml:"enabled"` VCPUs int `toml:"vcpus"` @@ -34,6 +42,11 @@ type Limits struct { Subnets int `toml:"subnets"` EIPs int `toml:"eips"` VolumesGiB int `toml:"volumes_gib"` + + TokensPerMonthEnabled bool `toml:"tokens_per_month_enabled"` + TokensPerMonth int64 `toml:"tokens_per_month"` + RequestsPerMinuteEnabled bool `toml:"requests_per_minute_enabled"` + RequestsPerMinute int `toml:"requests_per_minute"` } // Service enforces per-account quotas for one gateway. It holds the configured @@ -45,13 +58,24 @@ type Service struct { // quotas are disabled, in which case Exempt short-circuits every check // before the counter is touched. usage jetstream.KeyValue + + // bedrockUsage resolves the stream-fed Bedrock token counter + // (mulga-vmfng.7.6). Nil until SetBedrockUsage is called, in which case + // CheckBedrockTokens treats the dimension as unconfigured rather than + // erroring — a wiring gap must not block every Bedrock call. + bedrockUsage BedrockUsageReader + + // rpm holds the local, per-gateway, per-account token buckets enforcing + // the requests-per-minute dimension. Always non-nil so CheckBedrockRPM + // never needs a nil guard beyond Service itself. + rpm *rpmLimiter } // New constructs a quota Service from the configured limits and the gateway-owned // account-usage KV bucket. usage may be nil when quotas are disabled; Exempt then // short-circuits every check before the counter is read. func New(limits Limits, usage jetstream.KeyValue) *Service { - return &Service{limits: limits, usage: usage} + return &Service{limits: limits, usage: usage, rpm: newRPMLimiter()} } // Exempt returns true for the global/system account and whenever quotas are diff --git a/spinifex/services/awsgw/awsgw.go b/spinifex/services/awsgw/awsgw.go index d5ff2bd89..73486d4e6 100644 --- a/spinifex/services/awsgw/awsgw.go +++ b/spinifex/services/awsgw/awsgw.go @@ -329,6 +329,19 @@ func launchService(config *config.ClusterConfig) error { bedrockRecorder := gateway_bedrock.NewStreamRecorder(js, bedrockLoggingConfig) go gateway_bedrock.NewDeliveryConsumer(objStore, bedrockLoggingConfig).Run(janitorCtx, deliveryConsumer) + // Bedrock usage/cost metering: a second durable consumer on the same + // invocation stream (LimitsPolicy retention lets both see every message) + // updates per-account/model/period counters, deduping on RequestID so + // at-least-once redelivery never double-counts. bedrockPrices resolves + // KV price overrides over the catalog's in-tree defaults. + bedrockUsage := gateway_bedrock.NewUsageStore(js, len(config.Nodes)) + bedrockPrices := gateway_bedrock.NewPriceStore(js, len(config.Nodes)) + usageConsumer, err := gateway_bedrock.EnsureUsageConsumer(janitorCtx, js) + if err != nil { + return fmt.Errorf("bedrock: ensure usage metering consumer: %w", err) + } + go gateway_bedrock.NewUsageConsumer(bedrockUsage, bedrockPrices).Run(janitorCtx, usageConsumer) + gw := gateway.GatewayConfig{ Debug: nodeConfig.AWSGW.Debug, DisableLogging: false, @@ -379,6 +392,10 @@ func launchService(config *config.ClusterConfig) error { } gw.Quota = handlers_quota.New(quotaCfg, usageBucket) + // Bedrock token quota reads the stream-fed usage counters built above, + // independent of whether the standing-infra dimensions are enabled. + gw.Quota.SetBedrockUsage(bedrockUsage) + // Leader-locked vCPU reconcile: the only path that lowers the counter, // recomputing it from the running-plus-stopped sweep so out-of-band // terminations free quota. Started only when quotas are enabled so default-off @@ -391,6 +408,12 @@ func launchService(config *config.ClusterConfig) error { go runQuotaReconcile(janitorCtx, gw.Quota, natsConn, activeAccountIDs(iamService), expectedNodes) } + // Bedrock RPM enforcement is local and immediate (never gated on + // quotaCfg.Enabled); this loop only pushes a periodic observability + // snapshot to KV and is itself a no-op when the RPM dimension is + // disabled. + go gw.Quota.RunBedrockRPMSync(janitorCtx, js, len(config.Nodes)) + handler := gw.SetupRoutes() // Load TLS certificate