refactor: convert the bedrock guardrail and provisioned stores to kvstore (#897)

spinifex/gateway/bedrock/arn.go

@@ -191,11 +191,7 @@ func resolveInferenceTarget(ctx context.Context, accountID, modelID string, stor
 		return "", modelID, errors.New(awserrors.ErrorResourceNotFoundException)
 	}
 
-	kv, err := store.bucket(ctx)
-	if err != nil {
-		return "", modelID, err
-	}
-	rec, found, err := getProvisionedRecord(ctx, kv, parsed.AccountID, parsed.ID)
+	rec, found, err := store.get(ctx, parsed.AccountID, parsed.ID)
 	if err != nil {
 		return "", modelID, err
 	}

spinifex/gateway/bedrock/guardrail.go

@@ -2,13 +2,10 @@ package gateway_bedrock
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
-	"fmt"
 	"slices"
 	"strconv"
 	"strings"
-	"sync"
 	"time"
 	"uuid"
 
@@ -16,7 +13,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/bedrock"
 	"github.com/aws/aws-sdk-go/service/bedrockruntime"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
-	"github.com/mulgadc/spinifex/spinifex/kvutil"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/nats-io/nats.go/jetstream"
 )
 
@@ -85,40 +82,25 @@ type GuardrailRecord struct {
 // JetStream KV bucket, mirroring the ProvisionedStore/LoggingConfigStore
 // gateway-direct-KV pattern.
 type GuardrailStore struct {
-	js       jetstream.JetStream
-	replicas int
+	store *kvstore.Store[GuardrailRecord]
 	// region is baked in at construction, like ProvisionedStore.region, so
 	// ARN building/parsing needs no extra parameter threaded through the
 	// fixed-arity route table.
 	region string
-
-	mu sync.Mutex
-	kv jetstream.KeyValue
 }
 
 // NewGuardrailStore constructs a GuardrailStore over js, replicated across
 // replicas nodes, minting ARNs for region.
 func NewGuardrailStore(js jetstream.JetStream, replicas int, region string) *GuardrailStore {
-	return &GuardrailStore{js: js, replicas: replicas, region: region}
-}
-
-// bucket lazily opens (or creates) the bedrock-guardrails KV bucket, caching
-// the handle, mirroring ProvisionedStore.bucket.
-func (s *GuardrailStore) bucket(ctx context.Context) (jetstream.KeyValue, error) {
-	if s.js == nil {
-		return nil, errors.New("bedrock: guardrail store has no JetStream client configured")
+	return &GuardrailStore{
+		store: kvstore.New[GuardrailRecord](js, kvstore.Config{
+			Name:     bedrockGuardrailBucket,
+			History:  bedrockGuardrailHistory,
+			Replicas: replicas,
+			Missing:  "bedrock: guardrail store has no JetStream client configured",
+		}),
+		region: region,
 	}
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	if s.kv != nil {
-		return s.kv, nil
-	}
-	kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, s.js, bedrockGuardrailBucket, bedrockGuardrailHistory, s.replicas)
-	if err != nil {
-		return nil, err
-	}
-	s.kv = kv
-	return kv, nil
 }
 
 // guardrailKey scopes every stored record to its owning account, so List's
@@ -128,10 +110,10 @@ func guardrailKey(accountID, id string) string {
 	return accountID + "/" + id
 }
 
-// getGuardrailRecord reads accountID's record for id, or (zero, false, nil)
-// if it does not exist.
-func getGuardrailRecord(ctx context.Context, kv jetstream.KeyValue, accountID, id string) (GuardrailRecord, bool, error) {
-	rec, _, found, err := getGuardrailRecordRevision(ctx, kv, guardrailKey(accountID, id))
+// get reads accountID's record for id, or (zero, false, nil) if it does not
+// exist.
+func (s *GuardrailStore) get(ctx context.Context, accountID, id string) (GuardrailRecord, bool, error) {
+	rec, _, found, err := s.getRevision(ctx, guardrailKey(accountID, id))
 	return rec, found, err
 }
 
@@ -145,47 +127,33 @@ func errGuardrailNotFound(id, version string) error {
 	return awserrors.Errorf(awserrors.ErrorResourceNotFoundException, "guardrail %q version %q not found", id, version)
 }
 
-// getGuardrailRecordRevision is getGuardrailRecord with the KV revision
-// surfaced too, for a CAS write.
-func getGuardrailRecordRevision(ctx context.Context, kv jetstream.KeyValue, key string) (GuardrailRecord, uint64, bool, error) {
-	entry, err := kv.Get(ctx, key)
+// getRevision is get with the KV revision surfaced too, for a CAS write.
+func (s *GuardrailStore) getRevision(ctx context.Context, key string) (GuardrailRecord, uint64, bool, error) {
+	rec, rev, err := s.store.Get(ctx, key)
+	if errors.Is(err, kvstore.ErrNotFound) {
+		return GuardrailRecord{}, 0, false, nil
+	}
 	if err != nil {
-		if errors.Is(err, jetstream.ErrKeyNotFound) {
-			return GuardrailRecord{}, 0, false, nil
-		}
-		return GuardrailRecord{}, 0, false, fmt.Errorf("kv get %s: %w", key, err)
+		return GuardrailRecord{}, 0, false, err
 	}
-	var rec GuardrailRecord
-	if err := json.Unmarshal(entry.Value(), &rec); err != nil {
-		return GuardrailRecord{}, 0, false, fmt.Errorf("decode %s: %w", key, err)
-	}
-	return rec, entry.Revision(), true, nil
+	return *rec, rev, true, nil
 }
 
-// putGuardrailRecord writes rec as a brand-new key (Create only).
-func putGuardrailRecord(ctx context.Context, kv jetstream.KeyValue, rec GuardrailRecord) error {
-	data, err := json.Marshal(rec)
-	if err != nil {
-		return fmt.Errorf("encode guardrail %s: %w", rec.GuardrailID, err)
-	}
-	if _, err := kv.Put(ctx, guardrailKey(rec.AccountID, rec.GuardrailID), data); err != nil {
-		return fmt.Errorf("kv put guardrail %s: %w", rec.GuardrailID, err)
-	}
-	return nil
+// put writes rec under its own account-scoped key.
+func (s *GuardrailStore) put(ctx context.Context, rec GuardrailRecord) error {
+	return s.store.Set(ctx, guardrailKey(rec.AccountID, rec.GuardrailID), &rec)
 }
 
-// updateGuardrailRecord CAS-writes rec back over rev, the revision it was
-// read at, so two concurrent mutations of the same guardrail never silently
-// clobber one another.
-func updateGuardrailRecord(ctx context.Context, kv jetstream.KeyValue, key string, rec GuardrailRecord, rev uint64) error {
-	data, err := json.Marshal(rec)
-	if err != nil {
-		return fmt.Errorf("encode guardrail %s: %w", rec.GuardrailID, err)
+// update CAS-writes rec back over rev, the revision it was read at, so two
+// concurrent mutations of the same guardrail never silently clobber one
+// another. A lost race is reported as a retryable ConflictException.
+func (s *GuardrailStore) update(ctx context.Context, key string, rec GuardrailRecord, rev uint64) error {
+	err := s.store.CompareAndSet(ctx, key, &rec, rev)
+	if errors.Is(err, kvstore.ErrConflict) {
+		return awserrors.Errorf(awserrors.ErrorConflictException,
+			"bedrock: guardrail %s was modified concurrently; retry the request", rec.GuardrailID)
 	}
-	if _, err := kv.Update(ctx, key, data, rev); err != nil {
-		return fmt.Errorf("kv update guardrail %s: %w", rec.GuardrailID, err)
-	}
-	return nil
+	return err
 }
 
 // contentPolicyToSDK translates the stored *Config shape back to the
@@ -336,11 +304,6 @@ func CreateGuardrail(ctx context.Context, accountID string, store *GuardrailStor
 		return nil, errors.New(awserrors.ErrorValidationException)
 	}
 
-	kv, err := store.bucket(ctx)
-	if err != nil {
-		return nil, err
-	}
-
 	id := uuid.NewV4().String()
 	arn := FormatGuardrailARN(store.region, accountID, id)
 	now := time.Now().UTC()
@@ -365,7 +328,7 @@ func CreateGuardrail(ctx context.Context, accountID string, store *GuardrailStor
 		},
 	}
 
-	if err := putGuardrailRecord(ctx, kv, rec); err != nil {
+	if err := store.put(ctx, rec); err != nil {
 		return nil, err
 	}
 
@@ -390,11 +353,7 @@ func GetGuardrail(ctx context.Context, accountID string, store *GuardrailStore,
 		return nil, err
 	}
 
-	kv, err := store.bucket(ctx)
-	if err != nil {
-		return nil, err
-	}
-	rec, found, err := getGuardrailRecord(ctx, kv, accountID, id)
+	rec, found, err := store.get(ctx, accountID, id)
 	if err != nil {
 		return nil, err
 	}
@@ -418,37 +377,10 @@ func GetGuardrail(ctx context.Context, accountID string, store *GuardrailStore,
 // creation time (id as a deterministic tie-breaker), so one tenant never sees
 // another's and repeated calls return a stable order.
 func ListGuardrails(ctx context.Context, accountID string, store *GuardrailStore, _ *bedrock.ListGuardrailsInput) (*bedrock.ListGuardrailsOutput, error) {
-	kv, err := store.bucket(ctx)
+	recs, err := store.store.List(ctx, accountID+"/")
 	if err != nil {
 		return nil, err
 	}
-	keys, err := kv.Keys(ctx)
-	if err != nil {
-		if errors.Is(err, jetstream.ErrNoKeysFound) {
-			return &bedrock.ListGuardrailsOutput{}, nil
-		}
-		return nil, fmt.Errorf("kv keys for guardrails: %w", err)
-	}
-
-	prefix := accountID + "/"
-	var recs []GuardrailRecord
-	for _, key := range keys {
-		if !strings.HasPrefix(key, prefix) {
-			continue
-		}
-		entry, err := kv.Get(ctx, key)
-		if err != nil {
-			if errors.Is(err, jetstream.ErrKeyNotFound) {
-				continue
-			}
-			return nil, fmt.Errorf("kv get %s: %w", key, err)
-		}
-		var rec GuardrailRecord
-		if err := json.Unmarshal(entry.Value(), &rec); err != nil {
-			continue
-		}
-		recs = append(recs, rec)
-	}
 
 	slices.SortFunc(recs, func(a, b GuardrailRecord) int {
 		if c := a.CreatedAt.Compare(b.CreatedAt); c != 0 {
@@ -487,12 +419,8 @@ func UpdateGuardrail(ctx context.Context, accountID string, store *GuardrailStor
 		return nil, err
 	}
 
-	kv, err := store.bucket(ctx)
-	if err != nil {
-		return nil, err
-	}
 	key := guardrailKey(accountID, id)
-	rec, rev, found, err := getGuardrailRecordRevision(ctx, kv, key)
+	rec, rev, found, err := store.getRevision(ctx, key)
 	if err != nil {
 		return nil, err
 	}
@@ -511,7 +439,7 @@ func UpdateGuardrail(ctx context.Context, accountID string, store *GuardrailStor
 	rec.ContextualGroundingPolicy = input.ContextualGroundingPolicyConfig
 	rec.UpdatedAt = time.Now().UTC()
 
-	if err := updateGuardrailRecord(ctx, kv, key, rec, rev); err != nil {
+	if err := store.update(ctx, key, rec, rev); err != nil {
 		return nil, err
 	}
 
@@ -536,21 +464,17 @@ func DeleteGuardrail(ctx context.Context, accountID string, store *GuardrailStor
 		return nil, err
 	}
 
-	kv, err := store.bucket(ctx)
-	if err != nil {
-		return nil, err
-	}
 	key := guardrailKey(accountID, id)
 	version := aws.StringValue(input.GuardrailVersion)
 
 	if version == "" {
-		if err := kv.Delete(ctx, key); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
-			return nil, fmt.Errorf("kv delete guardrail %s: %w", id, err)
+		if err := store.store.Delete(ctx, key); err != nil {
+			return nil, err
 		}
 		return &bedrock.DeleteGuardrailOutput{}, nil
 	}
 
-	rec, rev, found, err := getGuardrailRecordRevision(ctx, kv, key)
+	rec, rev, found, err := store.getRevision(ctx, key)
 	if err != nil {
 		return nil, err
 	}
@@ -563,7 +487,7 @@ func DeleteGuardrail(ctx context.Context, accountID string, store *GuardrailStor
 	}
 
 	delete(rec.Versions, version)
-	if err := updateGuardrailRecord(ctx, kv, key, rec, rev); err != nil {
+	if err := store.update(ctx, key, rec, rev); err != nil {
 		return nil, err
 	}
 	return &bedrock.DeleteGuardrailOutput{}, nil
@@ -582,12 +506,8 @@ func CreateGuardrailVersion(ctx context.Context, accountID string, store *Guardr
 		return nil, err
 	}
 
-	kv, err := store.bucket(ctx)
-	if err != nil {
-		return nil, err
-	}
 	key := guardrailKey(accountID, id)
-	rec, rev, found, err := getGuardrailRecordRevision(ctx, kv, key)
+	rec, rev, found, err := store.getRevision(ctx, key)
 	if err != nil {
 		return nil, err
 	}
@@ -602,7 +522,7 @@ func CreateGuardrailVersion(ctx context.Context, accountID string, store *Guardr
 	}
 	rec.Versions[version] = GuardrailVersionSnapshot{Version: version, guardrailView: rec.guardrailView}
 
-	if err := updateGuardrailRecord(ctx, kv, key, rec, rev); err != nil {
+	if err := store.update(ctx, key, rec, rev); err != nil {
 		return nil, err
 	}
 
@@ -624,11 +544,7 @@ func ApplyGuardrail(ctx context.Context, accountID string, store *GuardrailStore
 		return nil, err
 	}
 
-	kv, err := store.bucket(ctx)
-	if err != nil {
-		return nil, err
-	}
-	rec, found, err := getGuardrailRecord(ctx, kv, accountID, id)
+	rec, found, err := store.get(ctx, accountID, id)
 	if err != nil {
 		return nil, err
 	}

spinifex/gateway/bedrock/guardrail_inference.go

@@ -23,11 +23,7 @@ func loadGuardrailView(ctx context.Context, store *GuardrailStore, accountID, id
 		return guardrailView{}, err
 	}
 
-	kv, err := store.bucket(ctx)
-	if err != nil {
-		return guardrailView{}, err
-	}
-	rec, found, err := getGuardrailRecord(ctx, kv, accountID, id)
+	rec, found, err := store.get(ctx, accountID, id)
 	if err != nil {
 		return guardrailView{}, err
 	}

spinifex/gateway/bedrock/guardrail_test.go

@@ -451,3 +451,50 @@ func TestGuardrailID_ResolvesBareOrARN(t *testing.T) {
 	_, err = GetGuardrail(ctx, grCallerAccount, store, &bedrock.GetGuardrailInput{GuardrailIdentifier: aws.String(foreignARN)})
 	require.Error(t, err)
 }
+
+// TestListGuardrails_RejectsAnUndecodableRecord asserts a corrupt record fails
+// the listing rather than silently omitting one of the account's guardrails. A
+// short list is a wrong answer to a question about active security controls.
+func TestListGuardrails_RejectsAnUndecodableRecord(t *testing.T) {
+	store := newGuardrailTestStore(t)
+	ctx := context.Background()
+
+	_, err := CreateGuardrail(ctx, grCallerAccount, store, createGuardrailInput("readable"))
+	require.NoError(t, err)
+
+	kv, err := store.store.KV(ctx)
+	require.NoError(t, err)
+	_, err = kv.Put(ctx, guardrailKey(grCallerAccount, "corrupt"), []byte("{not json"))
+	require.NoError(t, err)
+
+	_, err = ListGuardrails(ctx, grCallerAccount, store, nil)
+	require.Error(t, err)
+}
+
+// TestGuardrailStore_UpdateRejectsAStaleRevision covers the CAS guard: the
+// second writer at the same revision loses, and reports a retryable
+// ConflictException rather than clobbering the winner.
+func TestGuardrailStore_UpdateRejectsAStaleRevision(t *testing.T) {
+	store := newGuardrailTestStore(t)
+	ctx := context.Background()
+
+	out, err := CreateGuardrail(ctx, grCallerAccount, store, createGuardrailInput("cas"))
+	require.NoError(t, err)
+	key := guardrailKey(grCallerAccount, aws.StringValue(out.GuardrailId))
+
+	rec, stale, found, err := store.getRevision(ctx, key)
+	require.NoError(t, err)
+	require.True(t, found)
+
+	rec.Description = "first writer"
+	require.NoError(t, store.update(ctx, key, rec, stale))
+
+	rec.Description = "second writer"
+	err = store.update(ctx, key, rec, stale)
+	require.Error(t, err)
+	assert.True(t, awserrors.IsErrorCode(err, awserrors.ErrorConflictException))
+
+	got, _, _, err := store.getRevision(ctx, key)
+	require.NoError(t, err)
+	assert.Equal(t, "first writer", got.Description, "the loser must not have clobbered the winner")
+}

spinifex/gateway/bedrock/provisioned.go

@@ -2,19 +2,15 @@ package gateway_bedrock
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
-	"fmt"
 	"slices"
-	"strings"
-	"sync"
 	"time"
 	"uuid"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/bedrock"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
-	"github.com/mulgadc/spinifex/spinifex/kvutil"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/nats-io/nats.go/jetstream"
 )
 
@@ -77,42 +73,28 @@ type ProvisionedModelRecord struct {
 // each commitment, mirroring the LoggingConfigStore/ModelAccessStore
 // gateway-direct-KV pattern.
 type ProvisionedStore struct {
-	js       jetstream.JetStream
-	replicas int
+	store *kvstore.Store[ProvisionedModelRecord]
 	// region is baked in at construction (the gateway's own region never
 	// changes at runtime) so ARN building/parsing needs no extra parameter
 	// threaded through the fixed-arity route table.
 	region   string
 	endpoint EndpointProvisioner
-
-	mu sync.Mutex
-	kv jetstream.KeyValue
 }
 
 // NewProvisionedStore constructs a ProvisionedStore over js, replicated
 // across replicas nodes, driving endpoint launches/teardowns through
 // endpoint.
 func NewProvisionedStore(js jetstream.JetStream, replicas int, region string, endpoint EndpointProvisioner) *ProvisionedStore {
-	return &ProvisionedStore{js: js, replicas: replicas, region: region, endpoint: endpoint}
-}
-
-// bucket lazily opens (or creates) the bedrock-provisioned KV bucket, caching
-// the handle, mirroring LoggingConfigStore.bucket.
-func (s *ProvisionedStore) bucket(ctx context.Context) (jetstream.KeyValue, error) {
-	if s.js == nil {
-		return nil, errors.New("bedrock: provisioned store has no JetStream client configured")
+	return &ProvisionedStore{
+		store: kvstore.New[ProvisionedModelRecord](js, kvstore.Config{
+			Name:     bedrockProvisionedBucket,
+			History:  bedrockProvisionedHistory,
+			Replicas: replicas,
+			Missing:  "bedrock: provisioned store has no JetStream client configured",
+		}),
+		region:   region,
+		endpoint: endpoint,
 	}
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	if s.kv != nil {
-		return s.kv, nil
-	}
-	kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, s.js, bedrockProvisionedBucket, bedrockProvisionedHistory, s.replicas)
-	if err != nil {
-		return nil, err
-	}
-	s.kv = kv
-	return kv, nil
 }
 
 // provisionedKey scopes every stored record to its owning account, so List's
@@ -122,28 +104,36 @@ func provisionedKey(accountID, id string) string {
 	return accountID + "/" + id
 }
 
-// getProvisionedRecord reads accountID's record for id, or (zero, false, nil)
-// if it does not exist.
-func getProvisionedRecord(ctx context.Context, kv jetstream.KeyValue, accountID, id string) (ProvisionedModelRecord, bool, error) {
-	rec, _, found, err := getProvisionedRecordRevision(ctx, kv, provisionedKey(accountID, id))
+// get reads accountID's record for id, or (zero, false, nil) if it does not
+// exist.
+func (s *ProvisionedStore) get(ctx context.Context, accountID, id string) (ProvisionedModelRecord, bool, error) {
+	rec, _, found, err := s.getRevision(ctx, provisionedKey(accountID, id))
 	return rec, found, err
 }
 
-// getProvisionedRecordRevision is getProvisionedRecord with the KV revision
-// surfaced too, for Update's CAS write.
-func getProvisionedRecordRevision(ctx context.Context, kv jetstream.KeyValue, key string) (ProvisionedModelRecord, uint64, bool, error) {
-	entry, err := kv.Get(ctx, key)
+// getRevision is get with the KV revision surfaced too, for Update's CAS
+// write.
+func (s *ProvisionedStore) getRevision(ctx context.Context, key string) (ProvisionedModelRecord, uint64, bool, error) {
+	rec, rev, err := s.store.Get(ctx, key)
+	if errors.Is(err, kvstore.ErrNotFound) {
+		return ProvisionedModelRecord{}, 0, false, nil
+	}
 	if err != nil {
-		if errors.Is(err, jetstream.ErrKeyNotFound) {
-			return ProvisionedModelRecord{}, 0, false, nil
-		}
-		return ProvisionedModelRecord{}, 0, false, fmt.Errorf("kv get %s: %w", key, err)
+		return ProvisionedModelRecord{}, 0, false, err
 	}
-	var rec ProvisionedModelRecord
-	if err := json.Unmarshal(entry.Value(), &rec); err != nil {
-		return ProvisionedModelRecord{}, 0, false, fmt.Errorf("decode %s: %w", key, err)
+	return *rec, rev, true, nil
+}
+
+// update CAS-writes rec back over rev, the revision it was read at. A lost
+// race is reported as a retryable ConflictException rather than silently
+// clobbering the winner's write.
+func (s *ProvisionedStore) update(ctx context.Context, key string, rec ProvisionedModelRecord, rev uint64) error {
+	err := s.store.CompareAndSet(ctx, key, &rec, rev)
+	if errors.Is(err, kvstore.ErrConflict) {
+		return awserrors.Errorf(awserrors.ErrorConflictException,
+			"bedrock: provisioned throughput %s was modified concurrently; retry the request", rec.ARN)
 	}
-	return rec, entry.Revision(), true, nil
+	return err
 }
 
 // committedModelUnits sums ModelUnits across every commitment servingAccountID
@@ -156,35 +146,13 @@ func committedModelUnits(ctx context.Context, store *ProvisionedStore, servingAc
 	if servingAccountID == "" {
 		return 0, nil
 	}
-	kv, err := store.bucket(ctx)
+	recs, err := store.store.List(ctx, servingAccountID+"/")
 	if err != nil {
 		return 0, err
 	}
-	keys, err := kv.Keys(ctx)
-	if err != nil {
-		if errors.Is(err, jetstream.ErrNoKeysFound) {
-			return 0, nil
-		}
-		return 0, fmt.Errorf("kv keys for provisioned models: %w", err)
-	}
 
-	prefix := servingAccountID + "/"
 	var total int64
-	for _, key := range keys {
-		if !strings.HasPrefix(key, prefix) {
-			continue
-		}
-		entry, err := kv.Get(ctx, key)
-		if err != nil {
-			if errors.Is(err, jetstream.ErrKeyNotFound) {
-				continue
-			}
-			return 0, fmt.Errorf("kv get %s: %w", key, err)
-		}
-		var rec ProvisionedModelRecord
-		if err := json.Unmarshal(entry.Value(), &rec); err != nil {
-			continue
-		}
+	for _, rec := range recs {
 		if rec.ModelID == modelID {
 			total += rec.ModelUnits
 		}
@@ -241,8 +209,10 @@ func CreateProvisionedModelThroughput(ctx context.Context, accountID string, sto
 		return nil, errors.New(awserrors.ErrorResourceNotFoundException)
 	}
 
-	kv, err := store.bucket(ctx)
-	if err != nil {
+	// Resolve the bucket before pinning an endpoint, so an unreachable KV
+	// fails without launching a VM whose commitment record can never be
+	// written.
+	if _, err := store.store.KV(ctx); err != nil {
 		return nil, err
 	}
 
@@ -265,12 +235,8 @@ func CreateProvisionedModelThroughput(ctx context.Context, accountID string, sto
 		CreationTime:         now,
 		LastModifiedTime:     now,
 	}
-	data, err := json.Marshal(rec)
-	if err != nil {
-		return nil, fmt.Errorf("encode provisioned model %s: %w", id, err)
-	}
-	if _, err := kv.Put(ctx, provisionedKey(accountID, id), data); err != nil {
-		return nil, fmt.Errorf("kv put provisioned model %s: %w", id, err)
+	if err := store.store.Set(ctx, provisionedKey(accountID, id), &rec); err != nil {
+		return nil, err
 	}
 
 	return &bedrock.CreateProvisionedModelThroughputOutput{ProvisionedModelArn: aws.String(arn)}, nil
@@ -288,11 +254,7 @@ func GetProvisionedModelThroughput(ctx context.Context, accountID string, store
 		return nil, err
 	}
 
-	kv, err := store.bucket(ctx)
-	if err != nil {
-		return nil, err
-	}
-	rec, found, err := getProvisionedRecord(ctx, kv, accountID, id)
+	rec, found, err := store.get(ctx, accountID, id)
 	if err != nil {
 		return nil, err
 	}
@@ -327,37 +289,10 @@ func GetProvisionedModelThroughput(ctx context.Context, accountID string, store
 // Update), not live-derived — unlike Get, a list scan does not pay for one
 // endpoint describe per commitment.
 func ListProvisionedModelThroughputs(ctx context.Context, accountID string, store *ProvisionedStore, _ *bedrock.ListProvisionedModelThroughputsInput) (*bedrock.ListProvisionedModelThroughputsOutput, error) {
-	kv, err := store.bucket(ctx)
+	recs, err := store.store.List(ctx, accountID+"/")
 	if err != nil {
 		return nil, err
 	}
-	keys, err := kv.Keys(ctx)
-	if err != nil {
-		if errors.Is(err, jetstream.ErrNoKeysFound) {
-			return &bedrock.ListProvisionedModelThroughputsOutput{}, nil
-		}
-		return nil, fmt.Errorf("kv keys for provisioned models: %w", err)
-	}
-
-	prefix := accountID + "/"
-	var recs []ProvisionedModelRecord
-	for _, key := range keys {
-		if !strings.HasPrefix(key, prefix) {
-			continue
-		}
-		entry, err := kv.Get(ctx, key)
-		if err != nil {
-			if errors.Is(err, jetstream.ErrKeyNotFound) {
-				continue
-			}
-			return nil, fmt.Errorf("kv get %s: %w", key, err)
-		}
-		var rec ProvisionedModelRecord
-		if err := json.Unmarshal(entry.Value(), &rec); err != nil {
-			continue
-		}
-		recs = append(recs, rec)
-	}
 
 	slices.SortFunc(recs, func(a, b ProvisionedModelRecord) int {
 		return a.CreationTime.Compare(b.CreationTime)
@@ -396,12 +331,8 @@ func UpdateProvisionedModelThroughput(ctx context.Context, accountID string, sto
 		return nil, err
 	}
 
-	kv, err := store.bucket(ctx)
-	if err != nil {
-		return nil, err
-	}
 	key := provisionedKey(accountID, id)
-	rec, rev, found, err := getProvisionedRecordRevision(ctx, kv, key)
+	rec, rev, found, err := store.getRevision(ctx, key)
 	if err != nil {
 		return nil, err
 	}
@@ -418,12 +349,8 @@ func UpdateProvisionedModelThroughput(ctx context.Context, accountID string, sto
 	}
 	rec.LastModifiedTime = time.Now().UTC()
 
-	data, err := json.Marshal(rec)
-	if err != nil {
-		return nil, fmt.Errorf("encode provisioned model %s: %w", id, err)
-	}
-	if _, err := kv.Update(ctx, key, data, rev); err != nil {
-		return nil, fmt.Errorf("kv update provisioned model %s: %w", id, err)
+	if err := store.update(ctx, key, rec, rev); err != nil {
+		return nil, err
 	}
 	return &bedrock.UpdateProvisionedModelThroughputOutput{}, nil
 }
@@ -441,11 +368,7 @@ func DeleteProvisionedModelThroughput(ctx context.Context, accountID string, sto
 		return nil, err
 	}
 
-	kv, err := store.bucket(ctx)
-	if err != nil {
-		return nil, err
-	}
-	rec, found, err := getProvisionedRecord(ctx, kv, accountID, id)
+	rec, found, err := store.get(ctx, accountID, id)
 	if err != nil {
 		return nil, err
 	}
@@ -456,8 +379,8 @@ func DeleteProvisionedModelThroughput(ctx context.Context, accountID string, sto
 	if err := store.endpoint.DeletePinned(ctx, accountID, rec.ModelID); err != nil {
 		return nil, err
 	}
-	if err := kv.Delete(ctx, provisionedKey(accountID, id)); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
-		return nil, fmt.Errorf("kv delete provisioned model %s: %w", id, err)
+	if err := store.store.Delete(ctx, provisionedKey(accountID, id)); err != nil {
+		return nil, err
 	}
 	return &bedrock.DeleteProvisionedModelThroughputOutput{}, nil
 }

spinifex/gateway/bedrock/provisioned_test.go

@@ -341,3 +341,52 @@ func TestDeleteProvisionedModelThroughput_AbsentIsNoop(t *testing.T) {
 	require.NoError(t, err)
 	assert.Empty(t, stub.deleteCalls(), "an absent record must never reach the endpoint lifecycle")
 }
+
+// TestCommittedModelUnits_RejectsAnUndecodableRecord asserts a corrupt record
+// fails the sum rather than being skipped. Skipping undercounts committed
+// capacity, which admits traffic the commitments do not cover.
+func TestCommittedModelUnits_RejectsAnUndecodableRecord(t *testing.T) {
+	store := newProvisionedTestStore(t, newStubEndpointProvisioner())
+	ctx := context.Background()
+
+	_, err := CreateProvisionedModelThroughput(ctx, ptCallerAccount, store, createInput(selfHostTestModel, "readable", 2))
+	require.NoError(t, err)
+
+	kv, err := store.store.KV(ctx)
+	require.NoError(t, err)
+	_, err = kv.Put(ctx, provisionedKey(ptCallerAccount, "corrupt"), []byte("{not json"))
+	require.NoError(t, err)
+
+	_, err = committedModelUnits(ctx, store, ptCallerAccount, selfHostTestModel)
+	require.Error(t, err)
+}
+
+// TestProvisionedStore_UpdateRejectsAStaleRevision covers the CAS guard: the
+// second writer at the same revision loses, and reports a retryable
+// ConflictException rather than clobbering the winner.
+func TestProvisionedStore_UpdateRejectsAStaleRevision(t *testing.T) {
+	store := newProvisionedTestStore(t, newStubEndpointProvisioner())
+	ctx := context.Background()
+
+	out, err := CreateProvisionedModelThroughput(ctx, ptCallerAccount, store, createInput(selfHostTestModel, "cas", 1))
+	require.NoError(t, err)
+	parsed, err := ParseProvisionedModelARN(aws.StringValue(out.ProvisionedModelArn), ptTestRegion, ptCallerAccount)
+	require.NoError(t, err)
+	key := provisionedKey(ptCallerAccount, parsed.ID)
+
+	rec, stale, found, err := store.getRevision(ctx, key)
+	require.NoError(t, err)
+	require.True(t, found)
+
+	rec.ProvisionedModelName = "first-writer"
+	require.NoError(t, store.update(ctx, key, rec, stale))
+
+	rec.ProvisionedModelName = "second-writer"
+	err = store.update(ctx, key, rec, stale)
+	require.Error(t, err)
+	assert.True(t, awserrors.IsErrorCode(err, awserrors.ErrorConflictException))
+
+	got, _, _, err := store.getRevision(ctx, key)
+	require.NoError(t, err)
+	assert.Equal(t, "first-writer", got.ProvisionedModelName, "the loser must not have clobbered the winner")
+}