From f0b650a83b2101e8653eb034e9dfeef1cc84d578 Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Wed, 26 Aug 2026 23:38:33 +1000 Subject: [PATCH] refactor: convert the bedrock guardrail and provisioned stores to kvstore (#897) --- spinifex/gateway/bedrock/arn.go | 6 +- spinifex/gateway/bedrock/guardrail.go | 176 +++++------------ .../gateway/bedrock/guardrail_inference.go | 6 +- spinifex/gateway/bedrock/guardrail_test.go | 47 +++++ spinifex/gateway/bedrock/provisioned.go | 179 +++++------------- spinifex/gateway/bedrock/provisioned_test.go | 49 +++++ 6 files changed, 195 insertions(+), 268 deletions(-) diff --git a/spinifex/gateway/bedrock/arn.go b/spinifex/gateway/bedrock/arn.go index b3170a635..d1fcd3b47 100644 --- a/spinifex/gateway/bedrock/arn.go +++ b/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 } diff --git a/spinifex/gateway/bedrock/guardrail.go b/spinifex/gateway/bedrock/guardrail.go index 671cffcd4..fe00b257b 100644 --- a/spinifex/gateway/bedrock/guardrail.go +++ b/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 } diff --git a/spinifex/gateway/bedrock/guardrail_inference.go b/spinifex/gateway/bedrock/guardrail_inference.go index 695fb8adb..53058d14a 100644 --- a/spinifex/gateway/bedrock/guardrail_inference.go +++ b/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 } diff --git a/spinifex/gateway/bedrock/guardrail_test.go b/spinifex/gateway/bedrock/guardrail_test.go index 18cddfc2e..c1ab3517d 100644 --- a/spinifex/gateway/bedrock/guardrail_test.go +++ b/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") +} diff --git a/spinifex/gateway/bedrock/provisioned.go b/spinifex/gateway/bedrock/provisioned.go index 0cf8741c7..01c0262ff 100644 --- a/spinifex/gateway/bedrock/provisioned.go +++ b/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 } diff --git a/spinifex/gateway/bedrock/provisioned_test.go b/spinifex/gateway/bedrock/provisioned_test.go index cba9cbbbe..e4f9234ca 100644 --- a/spinifex/gateway/bedrock/provisioned_test.go +++ b/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") +}