feat: version jetstream buckets

spinifex/daemon/jetstream.go

@@ -9,6 +9,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/mulgadc/spinifex/spinifex/vm"
 	"github.com/nats-io/nats.go"
 )
@@ -26,6 +27,11 @@ const (
 	TerminatedInstanceBucket = "spinifex-terminated-instances"
 	// TerminatedInstancePrefix is the key prefix for terminated instances
 	TerminatedInstancePrefix = "terminated."
+
+	// Schema versions for daemon KV buckets
+	InstanceStateBucketVersion      = 1
+	ClusterStateBucketVersion       = 1
+	TerminatedInstanceBucketVersion = 1
 )
 
 // JetStreamManager manages JetStream KV store operations for instance state
@@ -77,6 +83,9 @@ func (m *JetStreamManager) InitKVBucket() error {
 	}
 
 	m.kv = kv
+	if err := utils.WriteVersion(kv, InstanceStateBucketVersion); err != nil {
+		return fmt.Errorf("write version to %s: %w", InstanceStateBucket, err)
+	}
 	return nil
 }
 
@@ -104,6 +113,9 @@ func (m *JetStreamManager) InitClusterStateBucket() error {
 	}
 
 	m.clusterKV = kv
+	if err := utils.WriteVersion(kv, ClusterStateBucketVersion); err != nil {
+		return fmt.Errorf("write version to %s: %w", ClusterStateBucket, err)
+	}
 	return nil
 }
 
@@ -132,6 +144,9 @@ func (m *JetStreamManager) InitTerminatedInstanceBucket() error {
 	}
 
 	m.terminatedKV = kv
+	if err := utils.WriteVersion(kv, TerminatedInstanceBucketVersion); err != nil {
+		return fmt.Errorf("write version to %s: %w", TerminatedInstanceBucket, err)
+	}
 	return nil
 }
 
@@ -157,7 +172,8 @@ func isStreamUnavailable(err error) bool {
 // recoverBucket attempts to reconnect to or re-create a KV bucket after the
 // underlying JetStream stream was lost during cluster formation.
 // Returns the recovered KV handle directly so callers avoid a racy re-read.
-func (m *JetStreamManager) recoverBucket(cfg *nats.KeyValueConfig, field *nats.KeyValue) (nats.KeyValue, error) {
+// When a bucket is recreated, the schema version is re-stamped.
+func (m *JetStreamManager) recoverBucket(cfg *nats.KeyValueConfig, field *nats.KeyValue, version int) (nats.KeyValue, error) {
 	m.kvMu.Lock()
 	defer m.kvMu.Unlock()
 
@@ -183,6 +199,10 @@ func (m *JetStreamManager) recoverBucket(cfg *nats.KeyValueConfig, field *nats.K
 		return nil, err
 	}
 
+	if err := utils.WriteVersion(kv, version); err != nil {
+		slog.Error("Failed to write version to recreated bucket", "bucket", cfg.Bucket, "err", err)
+	}
+
 	*field = kv
 	slog.Info("KV bucket recreated successfully", "bucket", cfg.Bucket)
 	return kv, nil
@@ -192,7 +212,7 @@ func (m *JetStreamManager) recoverKVBucket() (nats.KeyValue, error) {
 	return m.recoverBucket(&nats.KeyValueConfig{
 		Bucket:      InstanceStateBucket,
 		Description: "Spinifex instance state storage",
-	}, &m.kv)
+	}, &m.kv, InstanceStateBucketVersion)
 }
 
 func (m *JetStreamManager) recoverTerminatedKVBucket() (nats.KeyValue, error) {
@@ -200,7 +220,7 @@ func (m *JetStreamManager) recoverTerminatedKVBucket() (nats.KeyValue, error) {
 		Bucket:      TerminatedInstanceBucket,
 		Description: "Terminated instances (auto-expire after 1 hour)",
 		TTL:         1 * time.Hour,
-	}, &m.terminatedKV)
+	}, &m.terminatedKV, TerminatedInstanceBucketVersion)
 }
 
 // Heartbeat represents a daemon's periodic health status published to cluster KV.
@@ -662,6 +682,9 @@ func (m *JetStreamManager) ListStoppedInstances() ([]*vm.VM, error) {
 
 	var instances []*vm.VM
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(key, StoppedInstancePrefix) {
 			continue
 		}
@@ -751,6 +774,9 @@ func (m *JetStreamManager) ListTerminatedInstances() ([]*vm.VM, error) {
 
 	var instances []*vm.VM
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(key, TerminatedInstancePrefix) {
 			continue
 		}

spinifex/daemon/jetstream_test.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"testing"
 
+	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/mulgadc/spinifex/spinifex/vm"
 	"github.com/nats-io/nats.go"
 	"github.com/stretchr/testify/assert"
@@ -1208,3 +1209,58 @@ func TestJetStreamManager_ListTerminatedInstances_RecoverAfterStreamLost(t *test
 	require.NoError(t, err, "ListTerminatedInstances should recover after stream loss")
 	assert.Empty(t, instances)
 }
+
+func TestJetStreamManager_InitBuckets_WritesVersion(t *testing.T) {
+	nc, err := nats.Connect(sharedJSNATSURL)
+	require.NoError(t, err)
+	defer nc.Close()
+
+	jsm, err := NewJetStreamManager(nc, 1)
+	require.NoError(t, err)
+
+	require.NoError(t, jsm.InitKVBucket())
+	require.NoError(t, jsm.InitClusterStateBucket())
+	require.NoError(t, jsm.InitTerminatedInstanceBucket())
+
+	v, err := utils.ReadVersion(jsm.kv)
+	require.NoError(t, err)
+	assert.Equal(t, InstanceStateBucketVersion, v)
+
+	v, err = utils.ReadVersion(jsm.clusterKV)
+	require.NoError(t, err)
+	assert.Equal(t, ClusterStateBucketVersion, v)
+
+	v, err = utils.ReadVersion(jsm.terminatedKV)
+	require.NoError(t, err)
+	assert.Equal(t, TerminatedInstanceBucketVersion, v)
+}
+
+func TestJetStreamManager_ListStoppedInstances_SkipsVersionKey(t *testing.T) {
+	nc, err := nats.Connect(sharedJSNATSURL)
+	require.NoError(t, err)
+	defer nc.Close()
+
+	jsm, err := NewJetStreamManager(nc, 1)
+	require.NoError(t, err)
+	require.NoError(t, jsm.InitKVBucket())
+
+	// _version key is written by InitKVBucket; listing should not include it
+	instances, err := jsm.ListStoppedInstances()
+	require.NoError(t, err)
+	assert.Empty(t, instances)
+}
+
+func TestJetStreamManager_ListTerminatedInstances_SkipsVersionKey(t *testing.T) {
+	nc, err := nats.Connect(sharedJSNATSURL)
+	require.NoError(t, err)
+	defer nc.Close()
+
+	jsm, err := NewJetStreamManager(nc, 1)
+	require.NoError(t, err)
+	require.NoError(t, jsm.InitTerminatedInstanceBucket())
+
+	// _version key is written by InitTerminatedInstanceBucket; listing should not include it
+	instances, err := jsm.ListTerminatedInstances()
+	require.NoError(t, err)
+	assert.Empty(t, instances)
+}

spinifex/handlers/ec2/account/service_impl.go

@@ -14,9 +14,10 @@ import (
 )
 
 const (
-	KVBucketAccountSettings = "spinifex-ec2-account-settings"
-	KeyEbsEncryptionDefault = "ebs-encryption-default"
-	KeySerialConsoleAccess  = "serial-console-access"
+	KVBucketAccountSettings        = "spinifex-ec2-account-settings"
+	KVBucketAccountSettingsVersion = 1
+	KeyEbsEncryptionDefault        = "ebs-encryption-default"
+	KeySerialConsoleAccess         = "serial-console-access"
 )
 
 // AccountSettingsRecord represents stored account settings
@@ -43,6 +44,9 @@ func NewAccountSettingsServiceImplWithNATS(cfg *config.Config, natsConn *nats.Co
 	if err != nil {
 		return nil, fmt.Errorf("failed to create account settings KV bucket: %w", err)
 	}
+	if err := utils.WriteVersion(settingsKV, KVBucketAccountSettingsVersion); err != nil {
+		return nil, fmt.Errorf("write version to %s: %w", KVBucketAccountSettings, err)
+	}
 
 	slog.Info("Account settings service initialized with JetStream KV", "bucket", KVBucketAccountSettings)
 

spinifex/handlers/ec2/eigw/service_impl.go

@@ -20,7 +20,10 @@ import (
 // Ensure EgressOnlyIGWServiceImpl implements EgressOnlyIGWService
 var _ EgressOnlyIGWService = (*EgressOnlyIGWServiceImpl)(nil)
 
-const KVBucketEgressOnlyIGW = "spinifex-eigw"
+const (
+	KVBucketEgressOnlyIGW        = "spinifex-eigw"
+	KVBucketEgressOnlyIGWVersion = 1
+)
 
 // EgressOnlyIGWRecord represents a stored Egress-only Internet Gateway
 type EgressOnlyIGWRecord struct {
@@ -49,6 +52,9 @@ func NewEgressOnlyIGWServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketEgressOnlyIGW, err)
 	}
+	if err := utils.WriteVersion(eigwKV, KVBucketEgressOnlyIGWVersion); err != nil {
+		return nil, fmt.Errorf("write version to %s: %w", KVBucketEgressOnlyIGW, err)
+	}
 
 	// Get or create VPC KV bucket for cross-resource ownership validation
 	vpcKV, err := utils.GetOrCreateKVBucket(js, handlers_ec2_vpc.KVBucketVPCs, 10)
@@ -161,6 +167,9 @@ func (s *EgressOnlyIGWServiceImpl) DescribeEgressOnlyInternetGateways(input *ec2
 	}
 
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(key, prefix) {
 			continue
 		}

spinifex/handlers/ec2/igw/service_impl.go

@@ -21,7 +21,10 @@ import (
 // Ensure IGWServiceImpl implements IGWService
 var _ IGWService = (*IGWServiceImpl)(nil)
 
-const KVBucketIGW = "spinifex-igw"
+const (
+	KVBucketIGW        = "spinifex-igw"
+	KVBucketIGWVersion = 1
+)
 
 // IGWRecord represents a stored Internet Gateway
 type IGWRecord struct {
@@ -51,6 +54,9 @@ func NewIGWServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*IGWSer
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketIGW, err)
 	}
+	if err := utils.WriteVersion(igwKV, KVBucketIGWVersion); err != nil {
+		return nil, fmt.Errorf("write version to %s: %w", KVBucketIGW, err)
+	}
 
 	// Get or create VPC KV bucket for cross-resource ownership validation
 	vpcKV, err := utils.GetOrCreateKVBucket(js, handlers_ec2_vpc.KVBucketVPCs, 10)
@@ -157,6 +163,9 @@ func (s *IGWServiceImpl) DescribeInternetGateways(input *ec2.DescribeInternetGat
 	foundIDs := make(map[string]bool)
 
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(key, prefix) {
 			continue
 		}

spinifex/handlers/ec2/snapshot/service_impl.go

@@ -26,7 +26,10 @@ import (
 // Ensure SnapshotServiceImpl implements SnapshotService
 var _ SnapshotService = (*SnapshotServiceImpl)(nil)
 
-const KVBucketVolumeSnapshots = "spinifex-volume-snapshots"
+const (
+	KVBucketVolumeSnapshots        = "spinifex-volume-snapshots"
+	KVBucketVolumeSnapshotsVersion = 1
+)
 
 // SnapshotServiceImpl implements SnapshotService with S3-backed storage
 type SnapshotServiceImpl struct {
@@ -70,6 +73,9 @@ func NewSnapshotServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*S
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketVolumeSnapshots, err)
 	}
+	if err := utils.WriteVersion(kv, KVBucketVolumeSnapshotsVersion); err != nil {
+		return nil, nil, fmt.Errorf("write version to %s: %w", KVBucketVolumeSnapshots, err)
+	}
 
 	slog.Info("Snapshot service initialized with JetStream KV", "bucket", KVBucketVolumeSnapshots)
 

spinifex/handlers/ec2/vpc/eni.go

@@ -16,7 +16,8 @@ import (
 )
 
 const (
-	KVBucketENIs = "spinifex-vpc-enis"
+	KVBucketENIs        = "spinifex-vpc-enis"
+	KVBucketENIsVersion = 1
 )
 
 // ENIRecord represents a stored Elastic Network Interface
@@ -197,6 +198,9 @@ func (s *VPCServiceImpl) DescribeNetworkInterfaces(input *ec2.DescribeNetworkInt
 	}
 
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(key, prefix) {
 			continue
 		}

spinifex/handlers/ec2/vpc/ipam.go

@@ -13,7 +13,8 @@ import (
 )
 
 const (
-	KVBucketIPAM = "spinifex-vpc-ipam"
+	KVBucketIPAM        = "spinifex-vpc-ipam"
+	KVBucketIPAMVersion = 1
 )
 
 // IPAMRecord tracks allocated IPs for a subnet.
@@ -34,6 +35,9 @@ func NewIPAM(js nats.JetStreamContext) (*IPAM, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to create IPAM KV bucket: %w", err)
 	}
+	if err := utils.WriteVersion(kv, KVBucketIPAMVersion); err != nil {
+		return nil, fmt.Errorf("write version to %s: %w", KVBucketIPAM, err)
+	}
 	return &IPAM{kv: kv}, nil
 }
 

spinifex/handlers/ec2/vpc/service_impl.go

@@ -26,6 +26,10 @@ const (
 	KVBucketVNICounter = "spinifex-vpc-vni-counter"
 	vniCounterKey      = "counter"
 	vniStart           = 100 // Starting VNI value (avoid 0 and low numbers)
+
+	KVBucketVPCsVersion       = 1
+	KVBucketSubnetsVersion    = 1
+	KVBucketVNICounterVersion = 1
 )
 
 // VPCRecord represents a stored VPC
@@ -73,21 +77,33 @@ func NewVPCServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*VPCSer
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketVPCs, err)
 	}
+	if err := utils.WriteVersion(vpcKV, KVBucketVPCsVersion); err != nil {
+		return nil, fmt.Errorf("write version to %s: %w", KVBucketVPCs, err)
+	}
 
 	subnetKV, err := utils.GetOrCreateKVBucket(js, KVBucketSubnets, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketSubnets, err)
 	}
+	if err := utils.WriteVersion(subnetKV, KVBucketSubnetsVersion); err != nil {
+		return nil, fmt.Errorf("write version to %s: %w", KVBucketSubnets, err)
+	}
 
 	vniKV, err := utils.GetOrCreateKVBucket(js, KVBucketVNICounter, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketVNICounter, err)
 	}
+	if err := utils.WriteVersion(vniKV, KVBucketVNICounterVersion); err != nil {
+		return nil, fmt.Errorf("write version to %s: %w", KVBucketVNICounter, err)
+	}
 
 	eniKV, err := utils.GetOrCreateKVBucket(js, KVBucketENIs, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketENIs, err)
 	}
+	if err := utils.WriteVersion(eniKV, KVBucketENIsVersion); err != nil {
+		return nil, fmt.Errorf("write version to %s: %w", KVBucketENIs, err)
+	}
 
 	ipam, err := NewIPAM(js)
 	if err != nil {
@@ -225,6 +241,9 @@ func (s *VPCServiceImpl) DeleteVpc(input *ec2.DeleteVpcInput, accountID string)
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 	for _, k := range subnetKeys {
+		if k == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(k, prefix) {
 			continue
 		}
@@ -271,6 +290,9 @@ func (s *VPCServiceImpl) DescribeVpcs(input *ec2.DescribeVpcsInput, accountID st
 	}
 
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(key, prefix) {
 			continue
 		}
@@ -368,6 +390,9 @@ func (s *VPCServiceImpl) CreateSubnet(input *ec2.CreateSubnetInput, accountID st
 	}
 
 	for _, k := range subnetKeys {
+		if k == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(k, prefix) {
 			continue
 		}
@@ -499,6 +524,9 @@ func (s *VPCServiceImpl) DescribeSubnets(input *ec2.DescribeSubnetsInput, accoun
 	}
 
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(key, prefix) {
 			continue
 		}
@@ -634,6 +662,9 @@ func (s *VPCServiceImpl) EnsureDefaultVPC(accountID string) error {
 	}
 
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(key, prefix) {
 			continue
 		}
@@ -730,6 +761,9 @@ func (s *VPCServiceImpl) GetDefaultSubnet(accountID string) (*SubnetRecord, erro
 	}
 
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(key, prefix) {
 			continue
 		}

spinifex/handlers/iam/service_impl.go

@@ -27,6 +27,12 @@ const (
 	KVBucketAccounts       = "spinifex-accounts"
 	KVBucketAccountCounter = "spinifex-account-counter"
 
+	KVBucketUsersVersion          = 1
+	KVBucketAccessKeysVersion     = 1
+	KVBucketPoliciesVersion       = 1
+	KVBucketAccountsVersion       = 1
+	KVBucketAccountCounterVersion = 1
+
 	maxAccessKeysPerUser = 2
 )
 
@@ -63,26 +69,41 @@ func NewIAMServiceImpl(natsConn *nats.Conn, masterKey []byte, clusterSize int) (
 	if err != nil {
 		return nil, fmt.Errorf("init users bucket: %w", err)
 	}
+	if err := utils.WriteVersion(usersBucket, KVBucketUsersVersion); err != nil {
+		return nil, fmt.Errorf("write version to %s: %w", KVBucketUsers, err)
+	}
 
 	accessKeysBucket, err := getOrCreateBucket(js, KVBucketAccessKeys, 5, replicas)
 	if err != nil {
 		return nil, fmt.Errorf("init access keys bucket: %w", err)
 	}
+	if err := utils.WriteVersion(accessKeysBucket, KVBucketAccessKeysVersion); err != nil {
+		return nil, fmt.Errorf("write version to %s: %w", KVBucketAccessKeys, err)
+	}
 
 	policiesBucket, err := getOrCreateBucket(js, KVBucketPolicies, 10, replicas)
 	if err != nil {
 		return nil, fmt.Errorf("init policies bucket: %w", err)
 	}
+	if err := utils.WriteVersion(policiesBucket, KVBucketPoliciesVersion); err != nil {
+		return nil, fmt.Errorf("write version to %s: %w", KVBucketPolicies, err)
+	}
 
 	accountsBucket, err := getOrCreateBucket(js, KVBucketAccounts, 5, replicas)
 	if err != nil {
 		return nil, fmt.Errorf("init accounts bucket: %w", err)
 	}
+	if err := utils.WriteVersion(accountsBucket, KVBucketAccountsVersion); err != nil {
+		return nil, fmt.Errorf("write version to %s: %w", KVBucketAccounts, err)
+	}
 
 	accountCounterBucket, err := getOrCreateBucket(js, KVBucketAccountCounter, 5, replicas)
 	if err != nil {
 		return nil, fmt.Errorf("init account counter bucket: %w", err)
 	}
+	if err := utils.WriteVersion(accountCounterBucket, KVBucketAccountCounterVersion); err != nil {
+		return nil, fmt.Errorf("write version to %s: %w", KVBucketAccountCounter, err)
+	}
 
 	decrypter, err := NewDecrypter(masterKey)
 	if err != nil {
@@ -231,6 +252,9 @@ func (s *IAMServiceImpl) ListUsers(accountID string, input *iam.ListUsersInput)
 	keyPrefix := accountID + "."
 	var users []*iam.User
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(key, keyPrefix) {
 			continue
 		}
@@ -714,7 +738,12 @@ func (s *IAMServiceImpl) IsEmpty() (bool, error) {
 		}
 		return false, fmt.Errorf("check users bucket: %w", err)
 	}
-	return len(keys) == 0, nil
+	for _, key := range keys {
+		if key != utils.VersionKey {
+			return false, nil
+		}
+	}
+	return true, nil
 }
 
 // ---------------------------------------------------------------------------
@@ -826,6 +855,9 @@ func (s *IAMServiceImpl) ListAccounts() ([]*Account, error) {
 
 	accounts := make([]*Account, 0, len(keys))
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		entry, err := s.accountsBucket.Get(key)
 		if err != nil {
 			if errors.Is(err, nats.ErrKeyNotFound) {
@@ -984,6 +1016,9 @@ func (s *IAMServiceImpl) ListPolicies(accountID string, input *iam.ListPoliciesI
 	keyPrefix := accountID + "."
 	var policies []*iam.Policy
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(key, keyPrefix) {
 			continue
 		}
@@ -1229,6 +1264,9 @@ func (s *IAMServiceImpl) buildAttachmentCounts(accountID string) (map[string]int
 	counts := make(map[string]int64)
 	keyPrefix := accountID + "."
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(key, keyPrefix) {
 			continue
 		}
@@ -1267,6 +1305,9 @@ func (s *IAMServiceImpl) countPolicyAttachments(accountID, policyARN string) (in
 	keyPrefix := accountID + "."
 	var count int64
 	for _, key := range keys {
+		if key == utils.VersionKey {
+			continue
+		}
 		if !strings.HasPrefix(key, keyPrefix) {
 			continue
 		}

spinifex/utils/account.go

@@ -1,6 +1,54 @@
 package utils
 
-import "github.com/nats-io/nats.go"
+import (
+	"errors"
+	"fmt"
+	"strconv"
+
+	"github.com/nats-io/nats.go"
+)
+
+// VersionKey is the well-known KV key used to store a bucket's schema version.
+const VersionKey = "_version"
+
+// WriteVersion writes the schema version to a bucket.
+// Only writes if the key doesn't exist or the stored version is older.
+// Returns an error if the stored value is corrupt (non-integer).
+func WriteVersion(kv nats.KeyValue, version int) error {
+	entry, err := kv.Get(VersionKey)
+	if err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
+		return fmt.Errorf("read current version: %w", err)
+	}
+	if err == nil {
+		stored, parseErr := strconv.Atoi(string(entry.Value()))
+		if parseErr != nil {
+			return fmt.Errorf("corrupted _version key (raw=%q): %w", string(entry.Value()), parseErr)
+		}
+		if stored >= version {
+			return nil
+		}
+	}
+	_, err = kv.PutString(VersionKey, strconv.Itoa(version))
+	return err
+}
+
+// ReadVersion reads the schema version from a bucket.
+// Returns 0 if the key is not set. Returns an error on network failures
+// or corrupt values so callers can distinguish "not set" from "unreachable".
+func ReadVersion(kv nats.KeyValue) (int, error) {
+	entry, err := kv.Get(VersionKey)
+	if err != nil {
+		if errors.Is(err, nats.ErrKeyNotFound) {
+			return 0, nil
+		}
+		return 0, fmt.Errorf("read version: %w", err)
+	}
+	v, parseErr := strconv.Atoi(string(entry.Value()))
+	if parseErr != nil {
+		return 0, fmt.Errorf("corrupted _version key (raw=%q): %w", string(entry.Value()), parseErr)
+	}
+	return v, nil
+}
 
 // GlobalAccountID is the root/system account. Pre-Phase4 resources are owned by this account.
 const GlobalAccountID = "000000000000"

spinifex/utils/account_test.go

@@ -1,6 +1,15 @@
 package utils
 
-import "testing"
+import (
+	"os"
+	"testing"
+	"time"
+
+	"github.com/nats-io/nats-server/v2/server"
+	"github.com/nats-io/nats.go"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
 
 func TestAccountKey(t *testing.T) {
 	tests := []struct {
@@ -52,3 +61,122 @@ func TestGlobalAccountID(t *testing.T) {
 		t.Error("GlobalAccountID should be a valid account ID")
 	}
 }
+
+// startJSNATSServer starts an embedded JetStream-enabled NATS server for testing.
+func startJSNATSServer(t *testing.T) *server.Server {
+	t.Helper()
+	tmpDir, err := os.MkdirTemp("", "nats-js-test-*")
+	require.NoError(t, err)
+	t.Cleanup(func() { os.RemoveAll(tmpDir) })
+
+	ns, err := server.NewServer(&server.Options{
+		Host:      "127.0.0.1",
+		Port:      -1,
+		JetStream: true,
+		StoreDir:  tmpDir,
+		NoLog:     true,
+		NoSigs:    true,
+	})
+	require.NoError(t, err)
+	go ns.Start()
+	require.True(t, ns.ReadyForConnections(5*time.Second))
+	t.Cleanup(func() { ns.Shutdown() })
+	return ns
+}
+
+func TestWriteVersion_WritesKey(t *testing.T) {
+	ns := startJSNATSServer(t)
+	nc, err := nats.Connect(ns.ClientURL())
+	require.NoError(t, err)
+	defer nc.Close()
+
+	js, err := nc.JetStream()
+	require.NoError(t, err)
+
+	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{Bucket: "test-write-version"})
+	require.NoError(t, err)
+
+	err = WriteVersion(kv, 1)
+	require.NoError(t, err)
+
+	entry, err := kv.Get(VersionKey)
+	require.NoError(t, err)
+	assert.Equal(t, "1", string(entry.Value()))
+}
+
+func TestReadVersion_ReturnsZeroWhenUnset(t *testing.T) {
+	ns := startJSNATSServer(t)
+	nc, err := nats.Connect(ns.ClientURL())
+	require.NoError(t, err)
+	defer nc.Close()
+
+	js, err := nc.JetStream()
+	require.NoError(t, err)
+
+	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{Bucket: "test-read-unset"})
+	require.NoError(t, err)
+
+	v, err := ReadVersion(kv)
+	require.NoError(t, err)
+	assert.Equal(t, 0, v)
+}
+
+func TestWriteVersion_IdempotentOnSameVersion(t *testing.T) {
+	ns := startJSNATSServer(t)
+	nc, err := nats.Connect(ns.ClientURL())
+	require.NoError(t, err)
+	defer nc.Close()
+
+	js, err := nc.JetStream()
+	require.NoError(t, err)
+
+	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{Bucket: "test-idempotent"})
+	require.NoError(t, err)
+
+	require.NoError(t, WriteVersion(kv, 1))
+	require.NoError(t, WriteVersion(kv, 1))
+
+	v, err := ReadVersion(kv)
+	require.NoError(t, err)
+	assert.Equal(t, 1, v)
+}
+
+func TestWriteVersion_UpdatesOnHigherVersion(t *testing.T) {
+	ns := startJSNATSServer(t)
+	nc, err := nats.Connect(ns.ClientURL())
+	require.NoError(t, err)
+	defer nc.Close()
+
+	js, err := nc.JetStream()
+	require.NoError(t, err)
+
+	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{Bucket: "test-upgrade"})
+	require.NoError(t, err)
+
+	require.NoError(t, WriteVersion(kv, 1))
+	require.NoError(t, WriteVersion(kv, 2))
+
+	v, err := ReadVersion(kv)
+	require.NoError(t, err)
+	assert.Equal(t, 2, v)
+}
+
+func TestWriteVersion_NoOpOnLowerVersion(t *testing.T) {
+	ns := startJSNATSServer(t)
+	nc, err := nats.Connect(ns.ClientURL())
+	require.NoError(t, err)
+	defer nc.Close()
+
+	js, err := nc.JetStream()
+	require.NoError(t, err)
+
+	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{Bucket: "test-downgrade"})
+	require.NoError(t, err)
+
+	require.NoError(t, WriteVersion(kv, 3))
+	require.NoError(t, WriteVersion(kv, 1))
+
+	v, err := ReadVersion(kv)
+	require.NoError(t, err)
+	assert.Equal(t, 3, v)
+}