refactor: migrate spinifex off legacy nats.JetStreamContext onto nats.go/jetstream (#584)

.golangci.yml

@@ -26,6 +26,7 @@ linters:
     - errorlint # catches ==/!= on errors, type assertions on errors, and %v instead of %w — all break on wrapped errors
     - exptostd # detects golang.org/x/exp/ functions replaceable by stdlib
     - fatcontext # detects nested contexts in loops
+    - forbidigo # blocks the legacy JetStream/KV symbols
     - gocheckcompilerdirectives # validates go compiler directive comments (//go:)
     - godot # checks that declaration comments end in a period
     - goprintffuncname # checks that printf-like functions are named with f at the end
@@ -60,6 +61,16 @@ linters:
     embeddedstructfieldcheck:
       forbid-mutex: true
 
+    forbidigo:
+      analyze-types: true
+      forbid:
+        - pattern: '\.JetStream$'
+          pkg: '^github\.com/nats-io/nats\.go$'
+          msg: legacy nc.JetStream() is superseded; construct with jetstream.New(nc).
+        - pattern: '^nats\.(JetStreamContext|KeyValue.*|KeyWatcher|Err(KeyNotFound|KeyExists|BucketNotFound|NoKeysFound|StreamNotFound))$'
+          pkg: '^github\.com/nats-io/nats\.go$'
+          msg: legacy nats KV types/sentinels are superseded; use github.com/nats-io/nats.go/jetstream.
+
     depguard:
       rules:
         main:

cmd/spinifex/cmd/admin.go

@@ -2466,7 +2466,9 @@ func initIAMServiceFromConfig() (*handlers_iam.IAMServiceImpl, *config.ClusterCo
 		return nil, nil, nil, nil, fmt.Errorf("load master key: %w", err)
 	}
 
-	svc, err := handlers_iam.NewIAMServiceImpl(nc, masterKey, len(cfg.Nodes))
+	// Background: this runs at CLI top level, where there is no request to
+	// inherit a deadline from and the process exits after the one command.
+	svc, err := handlers_iam.NewIAMServiceImpl(context.Background(), nc, masterKey, len(cfg.Nodes))
 	if err != nil {
 		nc.Close()
 		return nil, nil, nil, nil, fmt.Errorf("init IAM service: %w", err)
@@ -2500,7 +2502,7 @@ func runAccountCreate(cmd *cobra.Command, args []string) {
 	// Create default VPC for the new account (belt-and-suspenders: daemon also
 	// does this via iam.account.created event, but daemon may not be running).
 	nodeConfig := cfg.Nodes[cfg.Node]
-	vpcSvc, vpcErr := handlers_ec2_vpc.NewVPCServiceImplWithNATS(&nodeConfig, nc)
+	vpcSvc, vpcErr := handlers_ec2_vpc.NewVPCServiceImplWithNATS(context.Background(), &nodeConfig, nc)
 	if vpcErr != nil {
 		fmt.Fprintf(os.Stderr, "Warning: could not create default VPC service: %v\n", vpcErr)
 	} else if _, vpcErr = vpcSvc.EnsureDefaultVPC(accountID); vpcErr != nil {

spinifex/daemon/daemon.go

@@ -57,6 +57,7 @@ import (
 	handlers_elbv2 "github.com/mulgadc/spinifex/spinifex/handlers/elbv2"
 	handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
 	"github.com/mulgadc/spinifex/spinifex/instancetypes"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/network/external"
 	"github.com/mulgadc/spinifex/spinifex/network/external/dhcp"
 	"github.com/mulgadc/spinifex/spinifex/network/host"
@@ -66,6 +67,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/mulgadc/spinifex/spinifex/vm"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // ResourceManager handles the allocation and tracking of system resources.
@@ -1296,8 +1298,8 @@ func (d *Daemon) startCluster() error {
 	}
 
 	// Remove the obsolete spinifex-dhcp-leases bucket (idempotent).
-	if js, jsErr := d.natsConn.JetStream(); jsErr == nil {
-		if err := utils.DeleteKVBucketIfExists(js, "spinifex-dhcp-leases"); err != nil {
+	if js, jsErr := jetstream.New(d.natsConn); jsErr == nil {
+		if err := kvutil.DeleteBucketIfExists(d.ctx, js, "spinifex-dhcp-leases"); err != nil {
 			slog.Warn("Failed to delete obsolete spinifex-dhcp-leases KV bucket", "err", err)
 		}
 	}
@@ -1333,10 +1335,10 @@ func (d *Daemon) startCluster() error {
 
 	type snapResult struct {
 		svc *handlers_ec2_snapshot.SnapshotServiceImpl
-		kv  nats.KeyValue
+		kv  jetstream.KeyValue
 	}
 	snap, err := initServiceWithRetry("snapshot service", func() (snapResult, error) {
-		svc, kv, err := handlers_ec2_snapshot.NewSnapshotServiceImplWithNATS(d.config, d.natsConn)
+		svc, kv, err := handlers_ec2_snapshot.NewSnapshotServiceImplWithNATS(d.ctx, d.config, d.natsConn)
 		return snapResult{svc, kv}, err
 	})
 	if err != nil {
@@ -1348,42 +1350,42 @@ func (d *Daemon) startCluster() error {
 	d.tagsService = handlers_ec2_tags.NewTagsServiceImpl(d.config)
 
 	d.eigwService, err = initServiceWithRetry("EIGW service", func() (*handlers_ec2_eigw.EgressOnlyIGWServiceImpl, error) {
-		return handlers_ec2_eigw.NewEgressOnlyIGWServiceImplWithNATS(d.config, d.natsConn)
+		return handlers_ec2_eigw.NewEgressOnlyIGWServiceImplWithNATS(d.ctx, d.config, d.natsConn)
 	})
 	if err != nil {
 		return fmt.Errorf("failed to initialize EIGW service: %w", err)
 	}
 
 	d.igwService, err = initServiceWithRetry("IGW service", func() (*handlers_ec2_igw.IGWServiceImpl, error) {
-		return handlers_ec2_igw.NewIGWServiceImplWithNATS(d.config, d.natsConn)
+		return handlers_ec2_igw.NewIGWServiceImplWithNATS(d.ctx, d.config, d.natsConn)
 	})
 	if err != nil {
 		return fmt.Errorf("failed to initialize IGW service: %w", err)
 	}
 
 	d.placementGroupService, err = initServiceWithRetry("placement group service", func() (*handlers_ec2_placementgroup.PlacementGroupServiceImpl, error) {
-		return handlers_ec2_placementgroup.NewPlacementGroupServiceImplWithNATS(d.config, d.natsConn)
+		return handlers_ec2_placementgroup.NewPlacementGroupServiceImplWithNATS(d.ctx, d.config, d.natsConn)
 	})
 	if err != nil {
 		return fmt.Errorf("failed to initialize placement group service: %w", err)
 	}
 
 	d.launchTemplateService, err = initServiceWithRetry("launch template service", func() (*handlers_ec2_launchtemplate.LaunchTemplateServiceImpl, error) {
-		return handlers_ec2_launchtemplate.NewLaunchTemplateServiceImplWithNATS(d.config, d.natsConn)
+		return handlers_ec2_launchtemplate.NewLaunchTemplateServiceImplWithNATS(d.ctx, d.config, d.natsConn)
 	})
 	if err != nil {
 		return fmt.Errorf("failed to initialize launch template service: %w", err)
 	}
 
 	d.spotInstanceService, err = initServiceWithRetry("spot instance service", func() (*handlers_ec2_spotinstance.SpotInstanceServiceImpl, error) {
-		return handlers_ec2_spotinstance.NewSpotInstanceServiceImplWithNATS(d.config, d.natsConn)
+		return handlers_ec2_spotinstance.NewSpotInstanceServiceImplWithNATS(d.ctx, d.config, d.natsConn)
 	})
 	if err != nil {
 		return fmt.Errorf("failed to initialize spot instance service: %w", err)
 	}
 
 	d.vpcService, err = initServiceWithRetry("VPC service", func() (*handlers_ec2_vpc.VPCServiceImpl, error) {
-		return handlers_ec2_vpc.NewVPCServiceImplWithNATS(d.config, d.natsConn)
+		return handlers_ec2_vpc.NewVPCServiceImplWithNATS(d.ctx, d.config, d.natsConn)
 	})
 	if err != nil {
 		return fmt.Errorf("failed to initialize VPC service: %w", err)
@@ -1394,7 +1396,7 @@ func (d *Daemon) startCluster() error {
 	d.vpcService.SetDefaultPublicIPMapping(d.hasPublicIPPools())
 
 	d.routeTableService, err = initServiceWithRetry("RouteTable service", func() (*handlers_ec2_routetable.RouteTableServiceImpl, error) {
-		return handlers_ec2_routetable.NewRouteTableServiceImplWithNATS(d.config, d.natsConn)
+		return handlers_ec2_routetable.NewRouteTableServiceImplWithNATS(d.ctx, d.config, d.natsConn)
 	})
 	if err != nil {
 		return fmt.Errorf("failed to initialize RouteTable service: %w", err)
@@ -1404,7 +1406,7 @@ func (d *Daemon) startCluster() error {
 	d.igwService.SetGatePublisher(d.routeTableService)
 
 	d.natGatewayService, err = initServiceWithRetry("NatGateway service", func() (*handlers_ec2_natgw.NatGatewayServiceImpl, error) {
-		return handlers_ec2_natgw.NewNatGatewayServiceImplWithNATS(d.natsConn)
+		return handlers_ec2_natgw.NewNatGatewayServiceImplWithNATS(d.ctx, d.natsConn)
 	})
 	if err != nil {
 		return fmt.Errorf("failed to initialize NatGateway service: %w", err)
@@ -1414,7 +1416,7 @@ func (d *Daemon) startCluster() error {
 	// mode with a public pool alongside the transit segment). The transit pool
 	// never enters IPAM — its addresses are gateway-LRP plumbing, not EIPs.
 	if d.hasPublicIPPools() {
-		js, jsErr := d.natsConn.JetStream()
+		js, jsErr := jetstream.New(d.natsConn)
 		if jsErr != nil {
 			slog.Warn("Failed to get JetStream for external IPAM", "err", jsErr)
 		} else {
@@ -1440,7 +1442,7 @@ func (d *Daemon) startCluster() error {
 					anyDHCP = true
 				}
 			}
-			d.externalIPAM, err = handlers_ec2_vpc.NewExternalIPAM(js, pools)
+			d.externalIPAM, err = handlers_ec2_vpc.NewExternalIPAM(d.ctx, js, pools)
 			if err != nil {
 				slog.Warn("Failed to initialize external IPAM", "err", err)
 			} else {
@@ -1457,7 +1459,7 @@ func (d *Daemon) startCluster() error {
 
 	// Initialize EIP service if external IPAM is available
 	if d.externalIPAM != nil && d.vpcService != nil {
-		eipSvc, eipErr := handlers_ec2_eip.NewEIPServiceImpl(d.natsConn, d.externalIPAM, d.vpcService)
+		eipSvc, eipErr := handlers_ec2_eip.NewEIPServiceImpl(d.ctx, d.natsConn, d.externalIPAM, d.vpcService)
 		if eipErr != nil {
 			slog.Warn("Failed to initialize EIP service", "err", eipErr)
 		} else {
@@ -1467,11 +1469,11 @@ func (d *Daemon) startCluster() error {
 
 		// Inject external IPAM + EIP KV into VPC service so DeleteNetworkInterface
 		// can release auto-assigned public IPs and NAT rules.
-		eipJS, eipJSErr := d.natsConn.JetStream()
+		eipJS, eipJSErr := jetstream.New(d.natsConn)
 		if eipJSErr != nil {
 			slog.Warn("Failed to get JetStream for VPC external IPAM injection", "err", eipJSErr)
 		} else {
-			eipKV, eipKVErr := utils.GetOrCreateKVBucket(eipJS, handlers_ec2_eip.KVBucketEIPs, 10)
+			eipKV, eipKVErr := kvutil.GetOrCreateBucket(d.ctx, eipJS, handlers_ec2_eip.KVBucketEIPs, 10)
 			if eipKVErr != nil {
 				slog.Warn("Failed to get EIP KV bucket for VPC service", "err", eipKVErr)
 			} else {
@@ -1495,7 +1497,7 @@ func (d *Daemon) startCluster() error {
 	}
 
 	d.accountService, err = initServiceWithRetry("account settings service", func() (*handlers_ec2_account.AccountSettingsServiceImpl, error) {
-		return handlers_ec2_account.NewAccountSettingsServiceImplWithNATS(d.config, d.natsConn)
+		return handlers_ec2_account.NewAccountSettingsServiceImplWithNATS(d.ctx, d.config, d.natsConn)
 	})
 	if err != nil {
 		return fmt.Errorf("failed to initialize account settings service: %w", err)
@@ -1547,9 +1549,9 @@ func (d *Daemon) startCluster() error {
 	// reaper. The scheduler is disabled (handlers still serve) when JetStream is
 	// unavailable.
 	d.ecsService = handlers_ecs.NewService(d.natsConn, d.config.Region, d.clusterConfig.AWS.InternalSuffix).WithDeps(d.buildECSServiceDeps())
-	if js, jsErr := d.natsConn.JetStream(); jsErr != nil {
+	if js, jsErr := jetstream.New(d.natsConn); jsErr != nil {
 		slog.Warn("ECS scheduler disabled: JetStream unavailable", "err", jsErr)
-	} else if _, lbErr := handlers_ecs.InitLeaderBucket(js); lbErr != nil {
+	} else if _, lbErr := handlers_ecs.InitLeaderBucket(d.ctx, js); lbErr != nil {
 		slog.Warn("ECS scheduler disabled: leader bucket init failed", "err", lbErr)
 	} else {
 		d.ecsScheduler = handlers_ecs.NewScheduler(d.natsConn, d.ecsService, d.node)
@@ -1564,7 +1566,7 @@ func (d *Daemon) startCluster() error {
 	}
 
 	d.acmService, err = initServiceWithRetry("ACM service", func() (*handlers_acm.ACMServiceImpl, error) {
-		return handlers_acm.NewACMServiceImplWithNATS(d.config, d.natsConn)
+		return handlers_acm.NewACMServiceImplWithNATS(d.ctx, d.config, d.natsConn)
 	})
 	if err != nil {
 		return fmt.Errorf("failed to initialize ACM service: %w", err)
@@ -1573,7 +1575,7 @@ func (d *Daemon) startCluster() error {
 	// ECR metadata service: owns per-account JetStream KV for repos, tags,
 	// manifest records and upload-state CAS. Disabled (gateway returns NATS
 	// timeouts) when JetStream is unavailable.
-	if js, jsErr := d.natsConn.JetStream(); jsErr != nil {
+	if js, jsErr := jetstream.New(d.natsConn); jsErr != nil {
 		slog.Warn("ECR metadata service disabled: JetStream unavailable", "err", jsErr)
 	} else {
 		d.ecrMetaService = handlers_ecr.NewKVMetaService(js)

spinifex/daemon/daemon_handlers_instance_tags_test.go

@@ -13,7 +13,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
 	"github.com/mulgadc/spinifex/spinifex/types"
 	"github.com/mulgadc/spinifex/spinifex/vm"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -54,11 +54,11 @@ func (m *memStoppedStore) DeleteStoppedInstance(id string) error {
 }
 
 // UpdateStoppedInstance mimics the real CAS semantics: a missing record
-// returns nats.ErrKeyNotFound instead of resurrecting it.
+// returns jetstream.ErrKeyNotFound instead of resurrecting it.
 func (m *memStoppedStore) UpdateStoppedInstance(id string, mutate func(*vm.VM)) (*vm.VM, error) {
 	v, ok := m.instances[id]
 	if !ok {
-		return nil, nats.ErrKeyNotFound
+		return nil, jetstream.ErrKeyNotFound
 	}
 	mutate(v)
 	return v, nil

spinifex/daemon/daemon_handlers_instance_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/mulgadc/spinifex/spinifex/vm"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -126,7 +127,7 @@ func (f *fakeStateStore) ClaimStoppedInstance(id string) (*vm.VM, error) {
 
 // UpdateStoppedInstance mimics the real CAS semantics for tests: mutate runs
 // under the store lock against the stored value, and a missing record
-// returns nats.ErrKeyNotFound (matching JetStreamManager's createIfAbsent=false
+// returns jetstream.ErrKeyNotFound (matching JetStreamManager's createIfAbsent=false
 // behavior) rather than resurrecting it.
 func (f *fakeStateStore) UpdateStoppedInstance(id string, mutate func(*vm.VM)) (*vm.VM, error) {
 	f.mu.Lock()
@@ -136,7 +137,7 @@ func (f *fakeStateStore) UpdateStoppedInstance(id string, mutate func(*vm.VM)) (
 	}
 	v, ok := f.stopped[id]
 	if !ok {
-		return nil, nats.ErrKeyNotFound
+		return nil, jetstream.ErrKeyNotFound
 	}
 	mutate(v)
 	return v, nil

spinifex/daemon/daemon_handlers_test.go

@@ -40,6 +40,7 @@ import (
 	"github.com/mulgadc/viperblock/viperblock"
 	"github.com/nats-io/nats-server/v2/server"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -132,7 +133,7 @@ func initAccountServiceForTest(t *testing.T, daemon *Daemon) {
 	require.NoError(t, err)
 	t.Cleanup(func() { nc.Close() })
 
-	svc, err := handlers_ec2_account.NewAccountSettingsServiceImplWithNATS(nil, nc)
+	svc, err := handlers_ec2_account.NewAccountSettingsServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 	daemon.accountService = svc
 }
@@ -2295,11 +2296,11 @@ func createVPCTestDaemon(t *testing.T) *Daemon {
 
 	testutil.StubVpcdSGResponder(t, nc)
 
-	vpcSvc, err := handlers_ec2_vpc.NewVPCServiceImplWithNATS(daemon.config, nc)
+	vpcSvc, err := handlers_ec2_vpc.NewVPCServiceImplWithNATS(t.Context(), daemon.config, nc)
 	require.NoError(t, err)
 	daemon.vpcService = vpcSvc
 
-	igwSvc, err := handlers_ec2_igw.NewIGWServiceImplWithNATS(daemon.config, nc)
+	igwSvc, err := handlers_ec2_igw.NewIGWServiceImplWithNATS(t.Context(), daemon.config, nc)
 	require.NoError(t, err)
 	daemon.igwService = igwSvc
 
@@ -2611,7 +2612,7 @@ func TestDelegateHandlers_EIGW(t *testing.T) {
 	require.NoError(t, err)
 	t.Cleanup(func() { nc.Close() })
 
-	eigwSvc, err := handlers_ec2_eigw.NewEgressOnlyIGWServiceImplWithNATS(daemon.config, nc)
+	eigwSvc, err := handlers_ec2_eigw.NewEgressOnlyIGWServiceImplWithNATS(t.Context(), daemon.config, nc)
 	require.NoError(t, err)
 	daemon.eigwService = eigwSvc
 
@@ -3563,14 +3564,14 @@ func TestDelegateHandlers_EIP(t *testing.T) {
 	require.NoError(t, err)
 	t.Cleanup(func() { nc.Close() })
 
-	js, err := nc.JetStream()
+	js, err := jetstream.New(nc)
 	require.NoError(t, err)
-	ipam, err := handlers_ec2_vpc.NewExternalIPAM(js, []external.ExternalPoolConfig{
+	ipam, err := handlers_ec2_vpc.NewExternalIPAM(t.Context(), js, []external.ExternalPoolConfig{
 		{Name: "test-pool", RangeStart: "192.168.100.2", RangeEnd: "192.168.100.254", Gateway: "192.168.100.1", PrefixLen: 24},
 	})
 	require.NoError(t, err)
 
-	eipSvc, err := handlers_ec2_eip.NewEIPServiceImpl(nc, ipam, daemon.vpcService)
+	eipSvc, err := handlers_ec2_eip.NewEIPServiceImpl(t.Context(), nc, ipam, daemon.vpcService)
 	require.NoError(t, err)
 	daemon.eipService = eipSvc
 
@@ -3756,7 +3757,7 @@ func TestDelegateHandlers_RouteTable(t *testing.T) {
 	require.NoError(t, err)
 	t.Cleanup(func() { nc.Close() })
 
-	rtbSvc, err := handlers_ec2_routetable.NewRouteTableServiceImplWithNATS(daemon.config, nc)
+	rtbSvc, err := handlers_ec2_routetable.NewRouteTableServiceImplWithNATS(t.Context(), daemon.config, nc)
 	require.NoError(t, err)
 	daemon.routeTableService = rtbSvc
 
@@ -3889,7 +3890,7 @@ func TestDelegateHandlers_PlacementGroup(t *testing.T) {
 	require.NoError(t, err)
 	t.Cleanup(func() { nc.Close() })
 
-	pgSvc, err := handlers_ec2_placementgroup.NewPlacementGroupServiceImplWithNATS(daemon.config, nc)
+	pgSvc, err := handlers_ec2_placementgroup.NewPlacementGroupServiceImplWithNATS(t.Context(), daemon.config, nc)
 	require.NoError(t, err)
 	daemon.placementGroupService = pgSvc
 

spinifex/daemon/daemon_mgmt_ip.go

@@ -8,7 +8,7 @@ import (
 	"sync"
 
 	"github.com/mulgadc/spinifex/spinifex/vm"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // mgmtIPAMKeyPrefix namespaces mgmt-IP allocation records within the shared
@@ -38,7 +38,7 @@ func updateMgmtIPAMWithRetry(jsManager *JetStreamManager, subnet string, mutate
 		if err == nil {
 			return record, nil
 		}
-		if !errors.Is(err, nats.ErrKeyExists) {
+		if !errors.Is(err, jetstream.ErrKeyExists) {
 			return nil, err
 		}
 		lastErr = err
@@ -247,7 +247,7 @@ func (a *MgmtIPAllocator) Release(instanceID string) {
 
 	// createIfAbsent=false: releasing an instance this node never allocated
 	// (or that no node has a record for) must not conjure an empty record
-	// into existence — nats.ErrKeyNotFound just means there was nothing to
+	// into existence — jetstream.ErrKeyNotFound just means there was nothing to
 	// release.
 	_, err := updateMgmtIPAMWithRetry(jsManager, subnet, func(r *MgmtIPRecord) {
 		for i, e := range r.Allocated {
@@ -257,7 +257,7 @@ func (a *MgmtIPAllocator) Release(instanceID string) {
 			}
 		}
 	}, false)
-	if err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
+	if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
 		slog.Warn("Failed to release mgmt IP in cluster KV; address remains reserved until reconciled",
 			"instance", instanceID, "err", err)
 	}

spinifex/daemon/daemon_mgmt_ip_test.go

@@ -1,6 +1,7 @@
 package daemon
 
 import (
+	"context"
 	"fmt"
 	"strings"
 	"sync"
@@ -38,7 +39,7 @@ func newTestMgmtJSM(t *testing.T) *JetStreamManager {
 func cleanupMgmtIPAM(t *testing.T, jsm *JetStreamManager, a *MgmtIPAllocator) {
 	t.Helper()
 	t.Cleanup(func() {
-		_ = jsm.clusterKV.Delete(mgmtIPAMKeyPrefix + a.subnet)
+		_ = jsm.clusterKV.Delete(context.Background(), mgmtIPAMKeyPrefix+a.subnet)
 	})
 }
 

spinifex/daemon/daemon_system_instance_test.go

@@ -21,6 +21,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/vm"
 	"github.com/nats-io/nats-server/v2/server"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -637,9 +638,9 @@ func TestRefreshSystemInstanceState_RewritesBlobs(t *testing.T) {
 	d.elbv2Service.GatewayURL = "https://10.0.0.1:9999"
 	d.elbv2Service.CACert = "-----BEGIN CERTIFICATE-----\nfake-ca\n-----END CERTIFICATE-----\n"
 
-	seedStore, err := handlers_elbv2.NewStore(nc)
+	seedStore, err := handlers_elbv2.NewStore(t.Context(), nc)
 	require.NoError(t, err)
-	require.NoError(t, seedStore.PutLoadBalancer(&handlers_elbv2.LoadBalancerRecord{
+	require.NoError(t, seedStore.PutLoadBalancer(t.Context(), &handlers_elbv2.LoadBalancerRecord{
 		LoadBalancerID: "lb-recover",
 		Name:           "recover-alb",
 		Scheme:         handlers_elbv2.SchemeInternal,
@@ -683,9 +684,9 @@ func TestRefreshSystemInstanceState_RewritesBlobs(t *testing.T) {
 func TestRefreshSystemInstanceState_WriteErrorPropagates(t *testing.T) {
 	d, nc, tmpDir := newRefreshTestDaemon(t)
 
-	seedStore, err := handlers_elbv2.NewStore(nc)
+	seedStore, err := handlers_elbv2.NewStore(t.Context(), nc)
 	require.NoError(t, err)
-	require.NoError(t, seedStore.PutLoadBalancer(&handlers_elbv2.LoadBalancerRecord{
+	require.NoError(t, seedStore.PutLoadBalancer(t.Context(), &handlers_elbv2.LoadBalancerRecord{
 		LoadBalancerID: "lb-werr",
 		Name:           "werr-alb",
 		Scheme:         handlers_elbv2.SchemeInternal,
@@ -747,9 +748,9 @@ func TestLaunchSystemInstance_NATFailureRollsBackPublicIP(t *testing.T) {
 	require.NoError(t, err)
 	t.Cleanup(func() { jsNC.Close() })
 
-	js, err := jsNC.JetStream()
+	js, err := jetstream.New(jsNC)
 	require.NoError(t, err)
-	ipam, err := handlers_ec2_vpc.NewExternalIPAM(js, []external.ExternalPoolConfig{
+	ipam, err := handlers_ec2_vpc.NewExternalIPAM(t.Context(), js, []external.ExternalPoolConfig{
 		{Name: "wan-test", RangeStart: "203.0.113.10", RangeEnd: "203.0.113.20", Gateway: "203.0.113.1", PrefixLen: 24},
 	})
 	require.NoError(t, err)
@@ -869,15 +870,15 @@ func TestReleaseSystemInstanceEIP_ReleasesEipServiceAllocation(t *testing.T) {
 	require.NoError(t, err)
 	t.Cleanup(func() { jsNC.Close() })
 
-	js, err := jsNC.JetStream()
+	js, err := jetstream.New(jsNC)
 	require.NoError(t, err)
-	ipam, err := handlers_ec2_vpc.NewExternalIPAM(js, []external.ExternalPoolConfig{
+	ipam, err := handlers_ec2_vpc.NewExternalIPAM(t.Context(), js, []external.ExternalPoolConfig{
 		{Name: "wan-test", RangeStart: "203.0.113.10", RangeEnd: "203.0.113.20", Gateway: "203.0.113.1", PrefixLen: 24},
 	})
 	require.NoError(t, err)
 	d.externalIPAM = ipam
 
-	eipSvc, err := handlers_ec2_eip.NewEIPServiceImpl(jsNC, ipam, d.vpcService)
+	eipSvc, err := handlers_ec2_eip.NewEIPServiceImpl(t.Context(), jsNC, ipam, d.vpcService)
 	require.NoError(t, err)
 	d.eipService = eipSvc
 

spinifex/daemon/daemon_test.go

@@ -54,6 +54,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/vm"
 	"github.com/mulgadc/viperblock/viperblock"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -1565,9 +1566,9 @@ func TestInstanceCleanerAdapter_DeleteVolumes_BootVolumeDeletedAfterAttachmentCl
 	jsConn, err := nats.Connect(sharedJSNATSURL)
 	require.NoError(t, err)
 	t.Cleanup(func() { jsConn.Close() })
-	js, err := jsConn.JetStream()
+	js, err := jetstream.New(jsConn)
 	require.NoError(t, err)
-	snapKV, err := js.CreateKeyValue(&nats.KeyValueConfig{
+	snapKV, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{
 		Bucket: "snap-kv-" + strings.ReplaceAll(t.Name(), "/", "-"),
 	})
 	require.NoError(t, err)
@@ -1659,9 +1660,9 @@ func TestTerminatedTeardownReaper_SelfHealsFailedVolumeTeardown(t *testing.T) {
 	daemon.stateStore = newStateStoreAdapter(jsManager)
 
 	// DeleteVolume's snapshot-reference guard requires a non-nil snapshotKV.
-	js, err := daemon.natsConn.JetStream()
+	js, err := jetstream.New(daemon.natsConn)
 	require.NoError(t, err)
-	snapKV, err := js.CreateKeyValue(&nats.KeyValueConfig{
+	snapKV, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{
 		Bucket: "snap-kv-" + strings.ReplaceAll(t.Name(), "/", "-"),
 	})
 	require.NoError(t, err)

spinifex/daemon/ecs_deps.go

@@ -46,7 +46,7 @@ func (d *Daemon) buildECSServiceDeps() handlers_ecs.Deps {
 	if d.clusterConfig != nil {
 		clusterSize = len(d.clusterConfig.Nodes)
 	}
-	if iamSvc, iamErr := handlers_iam.NewIAMServiceImpl(d.natsConn, masterKey, clusterSize); iamErr != nil {
+	if iamSvc, iamErr := handlers_iam.NewIAMServiceImpl(d.ctx, d.natsConn, masterKey, clusterSize); iamErr != nil {
 		slog.Warn("ECS: IAM service init failed; capacity provisioning disabled", "err", iamErr)
 	} else {
 		deps.IAM = iamSvc

spinifex/daemon/eks_deps.go

@@ -37,7 +37,7 @@ func (d *Daemon) systemRoleEnsurer() handlers_iam.SystemInstanceRoleEnsurer {
 	if d.clusterConfig != nil {
 		clusterSize = len(d.clusterConfig.Nodes)
 	}
-	iamSvc, iamErr := handlers_iam.NewIAMServiceImpl(d.natsConn, masterKey, clusterSize)
+	iamSvc, iamErr := handlers_iam.NewIAMServiceImpl(d.ctx, d.natsConn, masterKey, clusterSize)
 	if iamErr != nil {
 		slog.Warn("System role ensurer: IAM service init failed (retried on next launch); system VMs fall back to baked static creds",
 			"err", iamErr)

spinifex/daemon/heartbeat_test.go

@@ -124,7 +124,7 @@ func TestHeartbeatKVContract(t *testing.T) {
 
 	// Verify the KV key prefix and the JSON tag names directly. Reading back
 	// via ReadHeartbeat alone would also pass against a typo'd tag.
-	entry, err := jsm.clusterKV.Get("heartbeat." + h.Node)
+	entry, err := jsm.clusterKV.Get(t.Context(), "heartbeat."+h.Node)
 	require.NoError(t, err)
 
 	var raw map[string]any

spinifex/daemon/jetstream.go

@@ -1,6 +1,7 @@
 package daemon
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -13,6 +14,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/mulgadc/spinifex/spinifex/vm"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 const (
@@ -46,10 +48,10 @@ type KVSyncObserver interface {
 
 // JetStreamManager manages JetStream KV store operations for instance state.
 type JetStreamManager struct {
-	js           nats.JetStreamContext
-	kv           nats.KeyValue // spinifex-instance-state
-	clusterKV    nats.KeyValue // spinifex-cluster-state
-	terminatedKV nats.KeyValue // spinifex-terminated-instances
+	js           jetstream.JetStream
+	kv           jetstream.KeyValue // spinifex-instance-state
+	clusterKV    jetstream.KeyValue // spinifex-cluster-state
+	terminatedKV jetstream.KeyValue // spinifex-terminated-instances
 	replicas     int
 	kvMu         sync.Mutex // protects kv during recovery
 	obs          KVSyncObserver
@@ -68,14 +70,14 @@ func (m *JetStreamManager) KVHealthy() bool {
 	if m.js == nil {
 		return false
 	}
-	_, err := m.js.AccountInfo()
+	_, err := m.js.AccountInfo(context.Background())
 	return err == nil
 }
 
 // NewJetStreamManager creates a new JetStreamManager from a NATS connection.
 // replicas specifies the number of replicas for the KV bucket (typically matches cluster node count).
 func NewJetStreamManager(nc *nats.Conn, replicas int) (*JetStreamManager, error) {
-	js, err := nc.JetStream()
+	js, err := jetstream.New(nc)
 	if err != nil {
 		return nil, err
 	}
@@ -88,13 +90,14 @@ func NewJetStreamManager(nc *nats.Conn, replicas int) (*JetStreamManager, error)
 
 // InitKVBucket initializes the KV bucket, creating it if it doesn't exist.
 func (m *JetStreamManager) InitKVBucket() error {
+	ctx := context.Background()
 	// Try to get the existing bucket first
-	kv, err := m.js.KeyValue(InstanceStateBucket)
+	kv, err := m.js.KeyValue(ctx, InstanceStateBucket)
 	if err != nil {
-		if errors.Is(err, nats.ErrBucketNotFound) {
+		if errors.Is(err, jetstream.ErrBucketNotFound) {
 			// Bucket doesn't exist, create it
 			slog.Debug("Creating JetStream KV bucket", "bucket", InstanceStateBucket, "replicas", m.replicas)
-			kv, err = m.js.CreateKeyValue(&nats.KeyValueConfig{
+			kv, err = m.js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
 				Bucket:      InstanceStateBucket,
 				Description: "Spinifex instance state storage",
 				History:     1,          // Only keep latest value
@@ -111,7 +114,7 @@ func (m *JetStreamManager) InitKVBucket() error {
 	}
 
 	m.kv = kv
-	if err := migrate.DefaultRegistry.RunKV(InstanceStateBucket, kv, InstanceStateBucketVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, InstanceStateBucket, kv, InstanceStateBucketVersion); err != nil {
 		return fmt.Errorf("migrate %s: %w", InstanceStateBucket, err)
 	}
 	return nil
@@ -119,11 +122,12 @@ func (m *JetStreamManager) InitKVBucket() error {
 
 // InitClusterStateBucket initializes the cluster-state KV bucket, creating it if it doesn't exist.
 func (m *JetStreamManager) InitClusterStateBucket() error {
-	kv, err := m.js.KeyValue(ClusterStateBucket)
+	ctx := context.Background()
+	kv, err := m.js.KeyValue(ctx, ClusterStateBucket)
 	if err != nil {
-		if errors.Is(err, nats.ErrBucketNotFound) {
+		if errors.Is(err, jetstream.ErrBucketNotFound) {
 			slog.Debug("Creating JetStream KV bucket", "bucket", ClusterStateBucket, "replicas", m.replicas)
-			kv, err = m.js.CreateKeyValue(&nats.KeyValueConfig{
+			kv, err = m.js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
 				Bucket:      ClusterStateBucket,
 				Description: "Spinifex cluster state (heartbeats, shutdown markers, service maps)",
 				History:     1,
@@ -141,7 +145,7 @@ func (m *JetStreamManager) InitClusterStateBucket() error {
 	}
 
 	m.clusterKV = kv
-	if err := migrate.DefaultRegistry.RunKV(ClusterStateBucket, kv, ClusterStateBucketVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, ClusterStateBucket, kv, ClusterStateBucketVersion); err != nil {
 		return fmt.Errorf("migrate %s: %w", ClusterStateBucket, err)
 	}
 	return nil
@@ -150,11 +154,12 @@ func (m *JetStreamManager) InitClusterStateBucket() error {
 // InitTerminatedInstanceBucket initializes the terminated-instances KV bucket with a 1-hour TTL.
 // JetStream automatically purges keys after 1 hour, matching AWS behavior for terminated instances.
 func (m *JetStreamManager) InitTerminatedInstanceBucket() error {
-	kv, err := m.js.KeyValue(TerminatedInstanceBucket)
+	ctx := context.Background()
+	kv, err := m.js.KeyValue(ctx, TerminatedInstanceBucket)
 	if err != nil {
-		if errors.Is(err, nats.ErrBucketNotFound) {
+		if errors.Is(err, jetstream.ErrBucketNotFound) {
 			slog.Debug("Creating JetStream KV bucket", "bucket", TerminatedInstanceBucket, "replicas", m.replicas)
-			kv, err = m.js.CreateKeyValue(&nats.KeyValueConfig{
+			kv, err = m.js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
 				Bucket:      TerminatedInstanceBucket,
 				Description: "Terminated instances (auto-expire after 1 hour)",
 				History:     1,
@@ -172,7 +177,7 @@ func (m *JetStreamManager) InitTerminatedInstanceBucket() error {
 	}
 
 	m.terminatedKV = kv
-	if err := migrate.DefaultRegistry.RunKV(TerminatedInstanceBucket, kv, TerminatedInstanceBucketVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, TerminatedInstanceBucket, kv, TerminatedInstanceBucketVersion); err != nil {
 		return fmt.Errorf("migrate %s: %w", TerminatedInstanceBucket, err)
 	}
 	return nil
@@ -189,8 +194,8 @@ func isStreamUnavailable(err error) bool {
 	if err == nil {
 		return false
 	}
-	if errors.Is(err, nats.ErrStreamNotFound) ||
-		errors.Is(err, nats.ErrNoStreamResponse) ||
+	if errors.Is(err, jetstream.ErrStreamNotFound) ||
+		errors.Is(err, jetstream.ErrNoStreamResponse) ||
 		errors.Is(err, nats.ErrNoResponders) {
 		return true
 	}
@@ -201,19 +206,19 @@ func isStreamUnavailable(err error) bool {
 // underlying JetStream stream was lost during cluster formation.
 // Returns the recovered KV handle directly so callers avoid a racy re-read.
 // 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) {
+func (m *JetStreamManager) recoverBucket(ctx context.Context, cfg jetstream.KeyValueConfig, field *jetstream.KeyValue, version int) (jetstream.KeyValue, error) {
 	m.kvMu.Lock()
 	defer m.kvMu.Unlock()
 
 	// Try to reconnect to existing bucket first (another goroutine may have recovered it)
-	kv, err := m.js.KeyValue(cfg.Bucket)
+	kv, err := m.js.KeyValue(ctx, cfg.Bucket)
 	if err == nil {
 		*field = kv
 		slog.Info("Reconnected to KV bucket", "bucket", cfg.Bucket)
 		return kv, nil
 	}
 
-	if !errors.Is(err, nats.ErrBucketNotFound) && !isStreamUnavailable(err) {
+	if !errors.Is(err, jetstream.ErrBucketNotFound) && !isStreamUnavailable(err) {
 		return nil, err
 	}
 
@@ -221,13 +226,13 @@ func (m *JetStreamManager) recoverBucket(cfg *nats.KeyValueConfig, field *nats.K
 	slog.Warn("KV bucket stream lost, recreating", "bucket", cfg.Bucket, "replicas", m.replicas)
 	cfg.History = 1
 	cfg.Replicas = m.replicas
-	kv, err = m.js.CreateKeyValue(cfg)
+	kv, err = m.js.CreateKeyValue(ctx, cfg)
 	if err != nil {
 		slog.Error("Failed to recreate KV bucket", "bucket", cfg.Bucket, "err", err)
 		return nil, err
 	}
 
-	if err := migrate.DefaultRegistry.RunKV(cfg.Bucket, kv, version); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, cfg.Bucket, kv, version); err != nil {
 		slog.Error("Failed to run migrations on recreated bucket", "bucket", cfg.Bucket, "err", err)
 		return nil, fmt.Errorf("migrate recreated bucket %s: %w", cfg.Bucket, err)
 	}
@@ -237,15 +242,15 @@ func (m *JetStreamManager) recoverBucket(cfg *nats.KeyValueConfig, field *nats.K
 	return kv, nil
 }
 
-func (m *JetStreamManager) recoverKVBucket() (nats.KeyValue, error) {
-	return m.recoverBucket(&nats.KeyValueConfig{
+func (m *JetStreamManager) recoverKVBucket(ctx context.Context) (jetstream.KeyValue, error) {
+	return m.recoverBucket(ctx, jetstream.KeyValueConfig{
 		Bucket:      InstanceStateBucket,
 		Description: "Spinifex instance state storage",
 	}, &m.kv, InstanceStateBucketVersion)
 }
 
-func (m *JetStreamManager) recoverTerminatedKVBucket() (nats.KeyValue, error) {
-	return m.recoverBucket(&nats.KeyValueConfig{
+func (m *JetStreamManager) recoverTerminatedKVBucket(ctx context.Context) (jetstream.KeyValue, error) {
+	return m.recoverBucket(ctx, jetstream.KeyValueConfig{
 		Bucket:      TerminatedInstanceBucket,
 		Description: "Terminated instances (auto-expire after 1 hour)",
 		TTL:         1 * time.Hour,
@@ -260,10 +265,10 @@ const maxCASRetries = 5
 // isRevisionConflict reports whether err is the nats.go "wrong last
 // sequence" error that KeyValue.Update/Create return when the key's
 // revision no longer matches what the caller expected, i.e. a concurrent
-// writer won the race. nats.go surfaces this as nats.ErrKeyExists for both
-// Update (stale revision) and Create (key already present).
+// writer won the race. The jetstream package surfaces this as
+// jetstream.ErrKeyExists for both Update (stale revision) and Create (key already present).
 func isRevisionConflict(err error) bool {
-	return errors.Is(err, nats.ErrKeyExists)
+	return errors.Is(err, jetstream.ErrKeyExists)
 }
 
 // casUpdate performs an optimistic-concurrency read-modify-write against kv:
@@ -272,20 +277,20 @@ func isRevisionConflict(err error) bool {
 // Update/Create using the observed revision. A concurrent writer that changes
 // the key between Get and Update/Create is detected via a revision conflict
 // and retried, bounded by maxCASRetries.
-func casUpdate[T any](kv nats.KeyValue, key string, mutate func(*T), createIfAbsent bool) (*T, error) {
+func casUpdate[T any](ctx context.Context, kv jetstream.KeyValue, key string, mutate func(*T), createIfAbsent bool) (*T, error) {
 	var lastErr error
 	for range maxCASRetries {
 		var value T
 		var revision uint64
 
-		entry, err := kv.Get(key)
+		entry, err := kv.Get(ctx, key)
 		switch {
 		case err == nil:
 			if uerr := json.Unmarshal(entry.Value(), &value); uerr != nil {
 				return nil, uerr
 			}
 			revision = entry.Revision()
-		case errors.Is(err, nats.ErrKeyNotFound) && createIfAbsent:
+		case errors.Is(err, jetstream.ErrKeyNotFound) && createIfAbsent:
 			revision = 0
 		default:
 			return nil, err
@@ -299,9 +304,9 @@ func casUpdate[T any](kv nats.KeyValue, key string, mutate func(*T), createIfAbs
 		}
 
 		if revision == 0 {
-			_, err = kv.Create(key, data)
+			_, err = kv.Create(ctx, key, data)
 		} else {
-			_, err = kv.Update(key, data, revision)
+			_, err = kv.Update(ctx, key, data, revision)
 		}
 		if err == nil {
 			return &value, nil
@@ -322,25 +327,25 @@ func casUpdate[T any](kv nats.KeyValue, key string, mutate func(*T), createIfAbs
 // counterpart to casUpdate, for values (like vm.VM, which embeds a
 // sync.Mutex) that cannot be safely decoded into and copied out of a shared
 // struct value.
-func casPut(kv nats.KeyValue, key string, data []byte, createIfAbsent bool) error {
+func casPut(ctx context.Context, kv jetstream.KeyValue, key string, data []byte, createIfAbsent bool) error {
 	var lastErr error
 	for range maxCASRetries {
 		var revision uint64
 
-		entry, err := kv.Get(key)
+		entry, err := kv.Get(ctx, key)
 		switch {
 		case err == nil:
 			revision = entry.Revision()
-		case errors.Is(err, nats.ErrKeyNotFound) && createIfAbsent:
+		case errors.Is(err, jetstream.ErrKeyNotFound) && createIfAbsent:
 			revision = 0
 		default:
 			return err
 		}
 
 		if revision == 0 {
-			_, err = kv.Create(key, data)
+			_, err = kv.Create(ctx, key, data)
 		} else {
-			_, err = kv.Update(key, data, revision)
+			_, err = kv.Update(ctx, key, data, revision)
 		}
 		if err == nil {
 			return nil
@@ -363,12 +368,12 @@ func casPut(kv nats.KeyValue, key string, data []byte, createIfAbsent bool) erro
 // detected via a revision conflict and retried, bounded by maxCASRetries,
 // so a losing racer only fails outright when the key is truly gone
 // (notFound) or retries are exhausted.
-func casClaim[T any](kv nats.KeyValue, key string) (value *T, notFound bool, err error) {
+func casClaim[T any](ctx context.Context, kv jetstream.KeyValue, key string) (value *T, notFound bool, err error) {
 	var lastErr error
 	for range maxCASRetries {
-		entry, gerr := kv.Get(key)
+		entry, gerr := kv.Get(ctx, key)
 		if gerr != nil {
-			if errors.Is(gerr, nats.ErrKeyNotFound) {
+			if errors.Is(gerr, jetstream.ErrKeyNotFound) {
 				return nil, true, nil
 			}
 			return nil, false, gerr
@@ -379,7 +384,7 @@ func casClaim[T any](kv nats.KeyValue, key string) (value *T, notFound bool, err
 			return nil, false, uerr
 		}
 
-		if derr := kv.Delete(key, nats.LastRevision(entry.Revision())); derr != nil {
+		if derr := kv.Delete(ctx, key, jetstream.LastRevision(entry.Revision())); derr != nil {
 			if isRevisionConflict(derr) {
 				lastErr = derr
 				continue
@@ -422,7 +427,7 @@ func (m *JetStreamManager) WriteHeartbeat(h *Heartbeat) error {
 	if err != nil {
 		return err
 	}
-	_, err = m.clusterKV.Put("heartbeat."+h.Node, data)
+	_, err = m.clusterKV.Put(context.Background(), "heartbeat."+h.Node, data)
 	return err
 }
 
@@ -431,7 +436,7 @@ func (m *JetStreamManager) ReadHeartbeat(nodeID string) (*Heartbeat, error) {
 	if m.clusterKV == nil {
 		return nil, errors.New("cluster state KV not initialized")
 	}
-	entry, err := m.clusterKV.Get("heartbeat." + nodeID)
+	entry, err := m.clusterKV.Get(context.Background(), "heartbeat."+nodeID)
 	if err != nil {
 		return nil, err
 	}
@@ -461,7 +466,7 @@ func (m *JetStreamManager) WriteClusterShutdown(state *ClusterShutdownState) err
 	if m.clusterKV == nil {
 		return errors.New("cluster state KV not initialized")
 	}
-	_, err := casUpdate(m.clusterKV, "cluster.shutdown", func(s *ClusterShutdownState) {
+	_, err := casUpdate(context.Background(), m.clusterKV, "cluster.shutdown", func(s *ClusterShutdownState) {
 		*s = *state
 	}, true)
 	return err
@@ -470,12 +475,12 @@ func (m *JetStreamManager) WriteClusterShutdown(state *ClusterShutdownState) err
 // UpdateClusterShutdown atomically applies mutate to the current cluster
 // shutdown state (e.g. merging a node's ack into NodesAcked) and writes it
 // back with optimistic concurrency, retrying on a concurrent writer's
-// revision conflict. Returns nats.ErrKeyNotFound if no shutdown is in progress.
+// revision conflict. Returns jetstream.ErrKeyNotFound if no shutdown is in progress.
 func (m *JetStreamManager) UpdateClusterShutdown(mutate func(*ClusterShutdownState)) (*ClusterShutdownState, error) {
 	if m.clusterKV == nil {
 		return nil, errors.New("cluster state KV not initialized")
 	}
-	return casUpdate(m.clusterKV, "cluster.shutdown", mutate, false)
+	return casUpdate(context.Background(), m.clusterKV, "cluster.shutdown", mutate, false)
 }
 
 // UpdateMgmtIPAM atomically applies mutate to the mgmt-ipam record for
@@ -491,7 +496,7 @@ func (m *JetStreamManager) UpdateMgmtIPAM(subnet string, mutate func(*MgmtIPReco
 	if m.clusterKV == nil {
 		return nil, errors.New("cluster state KV not initialized")
 	}
-	return casUpdate(m.clusterKV, mgmtIPAMKeyPrefix+subnet, mutate, createIfAbsent)
+	return casUpdate(context.Background(), m.clusterKV, mgmtIPAMKeyPrefix+subnet, mutate, createIfAbsent)
 }
 
 // ReadClusterShutdown reads the cluster shutdown state from KV.
@@ -499,7 +504,7 @@ func (m *JetStreamManager) ReadClusterShutdown() (*ClusterShutdownState, error)
 	if m.clusterKV == nil {
 		return nil, errors.New("cluster state KV not initialized")
 	}
-	entry, err := m.clusterKV.Get("cluster.shutdown")
+	entry, err := m.clusterKV.Get(context.Background(), "cluster.shutdown")
 	if err != nil {
 		return nil, err
 	}
@@ -515,8 +520,8 @@ func (m *JetStreamManager) DeleteClusterShutdown() error {
 	if m.clusterKV == nil {
 		return errors.New("cluster state KV not initialized")
 	}
-	err := m.clusterKV.Delete("cluster.shutdown")
-	if err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
+	err := m.clusterKV.Delete(context.Background(), "cluster.shutdown")
+	if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
 		return err
 	}
 	return nil
@@ -534,7 +539,7 @@ func (m *JetStreamManager) WriteShutdownMarker(nodeID string) error {
 	if err != nil {
 		return fmt.Errorf("marshal shutdown marker: %w", err)
 	}
-	_, err = m.clusterKV.Put("shutdown."+nodeID, data)
+	_, err = m.clusterKV.Put(context.Background(), "shutdown."+nodeID, data)
 	return err
 }
 
@@ -543,8 +548,8 @@ func (m *JetStreamManager) ReadShutdownMarker(nodeID string) (bool, error) {
 	if m.clusterKV == nil {
 		return false, errors.New("cluster state KV not initialized")
 	}
-	_, err := m.clusterKV.Get("shutdown." + nodeID)
-	if errors.Is(err, nats.ErrKeyNotFound) {
+	_, err := m.clusterKV.Get(context.Background(), "shutdown."+nodeID)
+	if errors.Is(err, jetstream.ErrKeyNotFound) {
 		return false, nil
 	}
 	if err != nil {
@@ -558,8 +563,8 @@ func (m *JetStreamManager) DeleteShutdownMarker(nodeID string) error {
 	if m.clusterKV == nil {
 		return errors.New("cluster state KV not initialized")
 	}
-	err := m.clusterKV.Delete("shutdown." + nodeID)
-	if err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
+	err := m.clusterKV.Delete(context.Background(), "shutdown."+nodeID)
+	if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
 		return err
 	}
 	return nil
@@ -580,7 +585,7 @@ func (m *JetStreamManager) WriteServiceManifest(nodeID string, services []string
 	if err != nil {
 		return fmt.Errorf("marshal service manifest: %w", err)
 	}
-	_, err = m.clusterKV.Put("node."+nodeID+".services", data)
+	_, err = m.clusterKV.Put(context.Background(), "node."+nodeID+".services", data)
 	return err
 }
 
@@ -597,15 +602,15 @@ func (m *JetStreamManager) WriteState(nodeID string, vms map[string]*vm.VM) erro
 	}
 
 	key := InstanceStatePrefix + nodeID
-	_, err = m.kv.Put(key, jsonData)
+	_, err = m.kv.Put(context.Background(), key, jsonData)
 	if err != nil {
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "WriteState", "key", key, "err", err)
-			kv, recoverErr := m.recoverKVBucket()
+			kv, recoverErr := m.recoverKVBucket(context.Background())
 			if recoverErr != nil {
 				return err
 			}
-			if _, retryErr := kv.Put(key, jsonData); retryErr != nil {
+			if _, retryErr := kv.Put(context.Background(), key, jsonData); retryErr != nil {
 				return retryErr
 			}
 			slog.Debug("Wrote state to JetStream KV (after recovery)", "key", key, "instances", len(vms))
@@ -623,43 +628,32 @@ func (m *JetStreamManager) WriteState(nodeID string, vms map[string]*vm.VM) erro
 // blocks the caller past `timeout` and never returns an error. Used when the
 // local state file is the source of truth and KV is a best-effort cache; hot
 // paths marshal under a short-lived lock and commit lock-free.
-//
-// Note: kv.Put has no context API. On timeout, the in-flight Put goroutine
-// continues and completes (or fails) on its own. This leaks at most one
-// goroutine per write; under sustained partition the leak is bounded by
-// WriteState call cadence (per-state-transition, not a tight loop).
 func (m *JetStreamManager) WriteStateBytesBestEffort(nodeID string, jsonData []byte, timeout time.Duration) {
 	if m.kv == nil {
 		slog.Debug("KV bucket not initialized, skipping cluster sync", "node", nodeID)
 		return
 	}
 
-	key := InstanceStatePrefix + nodeID
-	done := make(chan error, 1)
-	go func() {
-		_, putErr := m.kv.Put(key, jsonData)
-		done <- putErr
-	}()
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
 
-	select {
-	case putErr := <-done:
-		if putErr != nil {
-			if m.obs != nil {
-				m.obs.RecordKVSyncFailure(InstanceStateBucket, putErr)
-			}
-			slog.Warn("KV sync failed (best-effort)", "key", key, "err", putErr)
+	key := InstanceStatePrefix + nodeID
+	_, err := m.kv.Put(ctx, key, jsonData)
+	if err != nil {
+		if m.obs != nil {
+			m.obs.RecordKVSyncFailure(InstanceStateBucket, err)
+		}
+		if errors.Is(err, context.DeadlineExceeded) {
+			slog.Warn("KV sync timed out (best-effort)", "key", key, "timeout", timeout)
 			return
 		}
-		if m.obs != nil {
-			m.obs.RecordKVSyncSuccess(InstanceStateBucket)
-		}
-		slog.Debug("Wrote state to KV (best-effort)", "key", key, "bytes", len(jsonData))
-	case <-time.After(timeout):
-		if m.obs != nil {
-			m.obs.RecordKVSyncFailure(InstanceStateBucket, fmt.Errorf("kv sync timeout after %s", timeout))
-		}
-		slog.Warn("KV sync timed out (best-effort)", "key", key, "timeout", timeout)
+		slog.Warn("KV sync failed (best-effort)", "key", key, "err", err)
+		return
 	}
+	if m.obs != nil {
+		m.obs.RecordKVSyncSuccess(InstanceStateBucket)
+	}
+	slog.Debug("Wrote state to KV (best-effort)", "key", key, "bytes", len(jsonData))
 }
 
 // marshalInstanceState produces the JSON wire form of vms.
@@ -680,22 +674,22 @@ func (m *JetStreamManager) LoadState(nodeID string) (map[string]*vm.VM, error) {
 	}
 
 	key := InstanceStatePrefix + nodeID
-	entry, err := m.kv.Get(key)
+	entry, err := m.kv.Get(context.Background(), key)
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			slog.Debug("No existing state in JetStream KV, returning empty state", "key", key)
 			return make(map[string]*vm.VM), nil
 		}
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "LoadState", "key", key, "err", err)
-			kv, recoverErr := m.recoverKVBucket()
+			kv, recoverErr := m.recoverKVBucket(context.Background())
 			if recoverErr != nil {
 				return nil, err
 			}
 			// Retry the read — if we reconnected, data may still exist
-			entry, err = kv.Get(key)
+			entry, err = kv.Get(context.Background(), key)
 			if err != nil {
-				if errors.Is(err, nats.ErrKeyNotFound) {
+				if errors.Is(err, jetstream.ErrKeyNotFound) {
 					slog.Warn("No state found after KV recovery", "key", key)
 					return make(map[string]*vm.VM), nil
 				}
@@ -728,16 +722,16 @@ func (m *JetStreamManager) DeleteState(nodeID string) error {
 	}
 
 	key := InstanceStatePrefix + nodeID
-	err := m.kv.Delete(key)
-	if err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
+	err := m.kv.Delete(context.Background(), key)
+	if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "DeleteState", "key", key, "err", err)
-			kv, recoverErr := m.recoverKVBucket()
+			kv, recoverErr := m.recoverKVBucket(context.Background())
 			if recoverErr != nil {
 				return err
 			}
 			// Retry — if we reconnected the key may still exist
-			if retryErr := kv.Delete(key); retryErr != nil && !errors.Is(retryErr, nats.ErrKeyNotFound) {
+			if retryErr := kv.Delete(context.Background(), key); retryErr != nil && !errors.Is(retryErr, jetstream.ErrKeyNotFound) {
 				return retryErr
 			}
 			return nil
@@ -760,14 +754,21 @@ func (m *JetStreamManager) UpdateReplicas(newReplicas int) error {
 
 	m.replicas = newReplicas
 
+	ctx := context.Background()
 	// Iterate all streams and update any KV-backed stream (prefixed "KV_")
 	updated := 0
-	for name := range m.js.StreamNames() {
+	lister := m.js.StreamNames(ctx)
+	for name := range lister.Name() {
 		if !strings.HasPrefix(name, "KV_") {
 			continue
 		}
 
-		info, err := m.js.StreamInfo(name)
+		stream, err := m.js.Stream(ctx, name)
+		if err != nil {
+			slog.Warn("Failed to open stream", "stream", name, "error", err)
+			continue
+		}
+		info, err := stream.Info(ctx)
 		if err != nil {
 			slog.Warn("Failed to get stream info", "stream", name, "error", err)
 			continue
@@ -779,7 +780,7 @@ func (m *JetStreamManager) UpdateReplicas(newReplicas int) error {
 
 		oldReplicas := info.Config.Replicas
 		info.Config.Replicas = newReplicas
-		if _, err := m.js.UpdateStream(&info.Config); err != nil {
+		if _, err := m.js.UpdateStream(ctx, info.Config); err != nil {
 			slog.Warn("Failed to update KV bucket replicas", "stream", name, "error", err)
 			continue
 		}
@@ -788,6 +789,9 @@ func (m *JetStreamManager) UpdateReplicas(newReplicas int) error {
 		slog.Info("Updated KV bucket replicas", "bucket", bucket, "oldReplicas", oldReplicas, "newReplicas", newReplicas)
 		updated++
 	}
+	if err := lister.Err(); err != nil {
+		return fmt.Errorf("list JetStream streams: %w", err)
+	}
 
 	if updated > 0 {
 		slog.Info("KV replication update complete", "bucketsUpdated", updated, "replicas", newReplicas)
@@ -813,15 +817,15 @@ func (m *JetStreamManager) WriteStoppedInstance(instanceID string, instance *vm.
 	}
 
 	key := StoppedInstancePrefix + instanceID
-	err = casPut(m.kv, key, jsonData, true)
+	err = casPut(context.Background(), m.kv, key, jsonData, true)
 	if err != nil {
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "WriteStoppedInstance", "key", key, "err", err)
-			kv, recoverErr := m.recoverKVBucket()
+			kv, recoverErr := m.recoverKVBucket(context.Background())
 			if recoverErr != nil {
 				return err
 			}
-			if retryErr := casPut(kv, key, jsonData, true); retryErr != nil {
+			if retryErr := casPut(context.Background(), kv, key, jsonData, true); retryErr != nil {
 				return retryErr
 			}
 			slog.Debug("Wrote stopped instance to JetStream KV (after recovery)", "key", key, "instanceId", instanceID)
@@ -842,21 +846,21 @@ func (m *JetStreamManager) LoadStoppedInstance(instanceID string) (*vm.VM, error
 	}
 
 	key := StoppedInstancePrefix + instanceID
-	entry, err := m.kv.Get(key)
+	entry, err := m.kv.Get(context.Background(), key)
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return nil, nil
 		}
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "LoadStoppedInstance", "key", key, "err", err)
-			kv, recoverErr := m.recoverKVBucket()
+			kv, recoverErr := m.recoverKVBucket(context.Background())
 			if recoverErr != nil {
 				return nil, err
 			}
 			// Retry — if we reconnected, data may still exist
-			entry, err = kv.Get(key)
+			entry, err = kv.Get(context.Background(), key)
 			if err != nil {
-				if errors.Is(err, nats.ErrKeyNotFound) {
+				if errors.Is(err, jetstream.ErrKeyNotFound) {
 					return nil, nil
 				}
 				return nil, err
@@ -883,16 +887,16 @@ func (m *JetStreamManager) DeleteStoppedInstance(instanceID string) error {
 	}
 
 	key := StoppedInstancePrefix + instanceID
-	err := m.kv.Delete(key)
-	if err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
+	err := m.kv.Delete(context.Background(), key)
+	if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "DeleteStoppedInstance", "key", key, "err", err)
-			kv, recoverErr := m.recoverKVBucket()
+			kv, recoverErr := m.recoverKVBucket(context.Background())
 			if recoverErr != nil {
 				return err
 			}
 			// Retry — if we reconnected the key may still exist
-			if retryErr := kv.Delete(key); retryErr != nil && !errors.Is(retryErr, nats.ErrKeyNotFound) {
+			if retryErr := kv.Delete(context.Background(), key); retryErr != nil && !errors.Is(retryErr, jetstream.ErrKeyNotFound) {
 				return retryErr
 			}
 			return nil
@@ -919,15 +923,15 @@ func (m *JetStreamManager) ClaimStoppedInstance(instanceID string) (*vm.VM, erro
 	}
 
 	key := StoppedInstancePrefix + instanceID
-	instance, notFound, err := casClaim[vm.VM](m.kv, key)
+	instance, notFound, err := casClaim[vm.VM](context.Background(), m.kv, key)
 	if err != nil {
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "ClaimStoppedInstance", "key", key, "err", err)
-			kv, recoverErr := m.recoverKVBucket()
+			kv, recoverErr := m.recoverKVBucket(context.Background())
 			if recoverErr != nil {
 				return nil, err
 			}
-			instance, notFound, err = casClaim[vm.VM](kv, key)
+			instance, notFound, err = casClaim[vm.VM](context.Background(), kv, key)
 			if err != nil {
 				return nil, err
 			}
@@ -947,25 +951,25 @@ func (m *JetStreamManager) ClaimStoppedInstance(instanceID string) (*vm.VM, erro
 // stopped record for instanceID and writes it back using optimistic
 // concurrency (CAS), retrying on a concurrent writer's revision conflict.
 // createIfAbsent is false: if a winning ClaimStoppedInstance deletes the
-// record mid-flight, the CAS write fails with nats.ErrKeyNotFound instead of
+// record mid-flight, the CAS write fails with jetstream.ErrKeyNotFound instead of
 // resurrecting it. Used by tag/attribute mutations that read-modify-write a
 // stopped instance so they cannot race a claim into recreating a stale
-// record. Returns nats.ErrKeyNotFound if no record exists.
+// record. Returns jetstream.ErrKeyNotFound if no record exists.
 func (m *JetStreamManager) UpdateStoppedInstance(instanceID string, mutate func(*vm.VM)) (*vm.VM, error) {
 	if m.kv == nil {
 		return nil, errors.New("KV bucket not initialized")
 	}
 
 	key := StoppedInstancePrefix + instanceID
-	updated, err := casUpdate(m.kv, key, mutate, false)
+	updated, err := casUpdate(context.Background(), m.kv, key, mutate, false)
 	if err != nil {
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "UpdateStoppedInstance", "key", key, "err", err)
-			kv, recoverErr := m.recoverKVBucket()
+			kv, recoverErr := m.recoverKVBucket(context.Background())
 			if recoverErr != nil {
 				return nil, err
 			}
-			return casUpdate(kv, key, mutate, false)
+			return casUpdate(context.Background(), kv, key, mutate, false)
 		}
 		return nil, err
 	}
@@ -978,21 +982,21 @@ func (m *JetStreamManager) ListStoppedInstances() ([]*vm.VM, error) {
 		return nil, errors.New("KV bucket not initialized")
 	}
 
-	keys, err := m.kv.Keys()
+	keys, err := m.kv.Keys(context.Background())
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, nil
 		}
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "ListStoppedInstances", "err", err)
-			kv, recoverErr := m.recoverKVBucket()
+			kv, recoverErr := m.recoverKVBucket(context.Background())
 			if recoverErr != nil {
 				return nil, err
 			}
 			// Retry — if we reconnected, data may still exist
-			keys, err = kv.Keys()
+			keys, err = kv.Keys(context.Background())
 			if err != nil {
-				if errors.Is(err, nats.ErrNoKeysFound) {
+				if errors.Is(err, jetstream.ErrNoKeysFound) {
 					return nil, nil
 				}
 				return nil, err
@@ -1012,9 +1016,9 @@ func (m *JetStreamManager) ListStoppedInstances() ([]*vm.VM, error) {
 			continue
 		}
 
-		entry, err := m.kv.Get(key)
+		entry, err := m.kv.Get(context.Background(), key)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
 			}
 			return nil, err
@@ -1051,15 +1055,15 @@ func (m *JetStreamManager) WriteTerminatedInstance(instanceID string, instance *
 	}
 
 	key := TerminatedInstancePrefix + instanceID
-	err = casPut(m.terminatedKV, key, jsonData, true)
+	err = casPut(context.Background(), m.terminatedKV, key, jsonData, true)
 	if err != nil {
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "WriteTerminatedInstance", "key", key, "err", err)
-			kv, recoverErr := m.recoverTerminatedKVBucket()
+			kv, recoverErr := m.recoverTerminatedKVBucket(context.Background())
 			if recoverErr != nil {
 				return err
 			}
-			if retryErr := casPut(kv, key, jsonData, true); retryErr != nil {
+			if retryErr := casPut(context.Background(), kv, key, jsonData, true); retryErr != nil {
 				return retryErr
 			}
 			slog.Debug("Wrote terminated instance to JetStream KV (after recovery)", "key", key, "instanceId", instanceID)
@@ -1077,22 +1081,22 @@ func (m *JetStreamManager) WriteTerminatedInstance(instanceID string, instance *
 // optimistic concurrency (CAS), retrying on a concurrent writer's revision
 // conflict. Used by the teardown reaper to merge per-dependent Teardown
 // progress without clobbering marks written by a concurrent update to the
-// same record. Returns nats.ErrKeyNotFound if no record exists yet.
+// same record. Returns jetstream.ErrKeyNotFound if no record exists yet.
 func (m *JetStreamManager) UpdateTerminatedInstance(instanceID string, mutate func(*vm.VM)) (*vm.VM, error) {
 	if m.terminatedKV == nil {
 		return nil, errors.New("terminated instance KV bucket not initialized")
 	}
 
 	key := TerminatedInstancePrefix + instanceID
-	updated, err := casUpdate(m.terminatedKV, key, mutate, false)
+	updated, err := casUpdate(context.Background(), m.terminatedKV, key, mutate, false)
 	if err != nil {
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "UpdateTerminatedInstance", "key", key, "err", err)
-			kv, recoverErr := m.recoverTerminatedKVBucket()
+			kv, recoverErr := m.recoverTerminatedKVBucket(context.Background())
 			if recoverErr != nil {
 				return nil, err
 			}
-			return casUpdate(kv, key, mutate, false)
+			return casUpdate(context.Background(), kv, key, mutate, false)
 		}
 		return nil, err
 	}
@@ -1105,20 +1109,20 @@ func (m *JetStreamManager) ListTerminatedInstances() ([]*vm.VM, error) {
 		return nil, errors.New("terminated instance KV bucket not initialized")
 	}
 
-	keys, err := m.terminatedKV.Keys()
+	keys, err := m.terminatedKV.Keys(context.Background())
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, nil
 		}
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "ListTerminatedInstances", "err", err)
-			kv, recoverErr := m.recoverTerminatedKVBucket()
+			kv, recoverErr := m.recoverTerminatedKVBucket(context.Background())
 			if recoverErr != nil {
 				return nil, err
 			}
-			keys, err = kv.Keys()
+			keys, err = kv.Keys(context.Background())
 			if err != nil {
-				if errors.Is(err, nats.ErrNoKeysFound) {
+				if errors.Is(err, jetstream.ErrNoKeysFound) {
 					return nil, nil
 				}
 				return nil, err
@@ -1137,9 +1141,9 @@ func (m *JetStreamManager) ListTerminatedInstances() ([]*vm.VM, error) {
 			continue
 		}
 
-		entry, err := m.terminatedKV.Get(key)
+		entry, err := m.terminatedKV.Get(context.Background(), key)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
 			}
 			return nil, err
@@ -1164,15 +1168,15 @@ func (m *JetStreamManager) DeleteTerminatedInstance(instanceID string) error {
 	}
 
 	key := TerminatedInstancePrefix + instanceID
-	err := m.terminatedKV.Delete(key)
-	if err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
+	err := m.terminatedKV.Delete(context.Background(), key)
+	if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "DeleteTerminatedInstance", "key", key, "err", err)
-			kv, recoverErr := m.recoverTerminatedKVBucket()
+			kv, recoverErr := m.recoverTerminatedKVBucket(context.Background())
 			if recoverErr != nil {
 				return err
 			}
-			if retryErr := kv.Delete(key); retryErr != nil && !errors.Is(retryErr, nats.ErrKeyNotFound) {
+			if retryErr := kv.Delete(context.Background(), key); retryErr != nil && !errors.Is(retryErr, jetstream.ErrKeyNotFound) {
 				return retryErr
 			}
 			return nil
@@ -1192,20 +1196,20 @@ func (m *JetStreamManager) LoadTerminatedInstance(instanceID string) (*vm.VM, er
 	}
 
 	key := TerminatedInstancePrefix + instanceID
-	entry, err := m.terminatedKV.Get(key)
+	entry, err := m.terminatedKV.Get(context.Background(), key)
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return nil, nil
 		}
 		if isStreamUnavailable(err) {
 			slog.Warn("KV stream unavailable, attempting recovery", "operation", "LoadTerminatedInstance", "key", key, "err", err)
-			kv, recoverErr := m.recoverTerminatedKVBucket()
+			kv, recoverErr := m.recoverTerminatedKVBucket(context.Background())
 			if recoverErr != nil {
 				return nil, err
 			}
-			entry, err = kv.Get(key)
+			entry, err = kv.Get(context.Background(), key)
 			if err != nil {
-				if errors.Is(err, nats.ErrKeyNotFound) {
+				if errors.Is(err, jetstream.ErrKeyNotFound) {
 					return nil, nil
 				}
 				return nil, err

spinifex/daemon/jetstream_test.go

@@ -7,9 +7,10 @@ import (
 	"testing"
 	"time"
 
-	"github.com/mulgadc/spinifex/spinifex/utils"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/vm"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -343,8 +344,10 @@ func TestJetStreamManager_UpdateReplicas(t *testing.T) {
 	require.NoError(t, err)
 
 	// Verify initial replica count
-	js, _ := nc.JetStream()
-	streamInfo, err := js.StreamInfo("KV_" + InstanceStateBucket)
+	js, _ := jetstream.New(nc)
+	stream, err := js.Stream(t.Context(), "KV_"+InstanceStateBucket)
+	require.NoError(t, err)
+	streamInfo, err := stream.Info(t.Context())
 	require.NoError(t, err)
 	assert.Equal(t, 1, streamInfo.Config.Replicas, "Should start with 1 replica")
 
@@ -598,7 +601,7 @@ func TestJetStreamManager_UpdateStoppedInstance_ConflictRetry(t *testing.T) {
 }
 
 // TestJetStreamManager_UpdateStoppedInstance_NotFound verifies
-// UpdateStoppedInstance surfaces nats.ErrKeyNotFound rather than silently
+// UpdateStoppedInstance surfaces jetstream.ErrKeyNotFound rather than silently
 // creating a record when nothing exists to update.
 func TestJetStreamManager_UpdateStoppedInstance_NotFound(t *testing.T) {
 	nc, err := nats.Connect(sharedJSNATSURL)
@@ -610,7 +613,7 @@ func TestJetStreamManager_UpdateStoppedInstance_NotFound(t *testing.T) {
 	require.NoError(t, jsm.InitKVBucket())
 
 	_, err = jsm.UpdateStoppedInstance("i-does-not-exist", func(v *vm.VM) {})
-	assert.ErrorIs(t, err, nats.ErrKeyNotFound)
+	assert.ErrorIs(t, err, jetstream.ErrKeyNotFound)
 }
 
 // TestJetStreamManager_UpdateStoppedInstance_NoResurrectAfterClaim is the
@@ -646,7 +649,7 @@ func TestJetStreamManager_UpdateStoppedInstance_NoResurrectAfterClaim(t *testing
 	_, err = jsm.UpdateStoppedInstance(testVM.ID, func(v *vm.VM) {
 		v.InstanceType = "should-not-land"
 	})
-	assert.ErrorIs(t, err, nats.ErrKeyNotFound, "a claim that deleted the record must not be resurrected by a losing racer's update")
+	assert.ErrorIs(t, err, jetstream.ErrKeyNotFound, "a claim that deleted the record must not be resurrected by a losing racer's update")
 
 	stillGone, err := jsm.LoadStoppedInstance(testVM.ID)
 	require.NoError(t, err)
@@ -776,7 +779,7 @@ func TestJetStreamManager_WriteServiceManifest(t *testing.T) {
 	require.NoError(t, err)
 
 	// Read the KV entry directly and verify JSON contents
-	entry, err := jsm.clusterKV.Get("node.test-node-svc.services")
+	entry, err := jsm.clusterKV.Get(t.Context(), "node.test-node-svc.services")
 	require.NoError(t, err)
 
 	var manifest map[string]any
@@ -811,7 +814,7 @@ func TestJetStreamManager_WriteServiceManifest_EmptyServices(t *testing.T) {
 	err = jsm.WriteServiceManifest("empty-svc-node", []string{}, "10.0.0.2:4222", "10.0.0.2:8443")
 	require.NoError(t, err)
 
-	entry, err := jsm.clusterKV.Get("node.empty-svc-node.services")
+	entry, err := jsm.clusterKV.Get(t.Context(), "node.empty-svc-node.services")
 	require.NoError(t, err)
 
 	var manifest map[string]any
@@ -847,9 +850,9 @@ func TestJetStreamManager_WriteServiceManifest_ClusterKVNotInitialized(t *testin
 // deleteInstanceStateBucket deletes the underlying JetStream stream for the instance-state KV bucket.
 func deleteInstanceStateBucket(t *testing.T, nc *nats.Conn) {
 	t.Helper()
-	js, err := nc.JetStream()
+	js, err := jetstream.New(nc)
 	require.NoError(t, err)
-	err = js.DeleteStream("KV_" + InstanceStateBucket)
+	err = js.DeleteStream(t.Context(), "KV_"+InstanceStateBucket)
 	require.NoError(t, err)
 }
 
@@ -1006,7 +1009,7 @@ func swapToNonJSContext(t *testing.T, jsm *JetStreamManager) {
 	nc, err := nats.Connect(sharedNATSURL)
 	require.NoError(t, err)
 	t.Cleanup(func() { nc.Close() })
-	js, err := nc.JetStream()
+	js, err := jetstream.New(nc)
 	require.NoError(t, err)
 	jsm.js = js
 }
@@ -1137,12 +1140,12 @@ func TestJetStreamManager_ListStoppedInstances_RecoveryFailure(t *testing.T) {
 
 func TestIsStreamUnavailable(t *testing.T) {
 	assert.False(t, isStreamUnavailable(nil))
-	assert.True(t, isStreamUnavailable(nats.ErrStreamNotFound))
-	assert.True(t, isStreamUnavailable(nats.ErrNoStreamResponse))
+	assert.True(t, isStreamUnavailable(jetstream.ErrStreamNotFound))
+	assert.True(t, isStreamUnavailable(jetstream.ErrNoStreamResponse))
 	assert.True(t, isStreamUnavailable(nats.ErrNoResponders))
 	assert.True(t, isStreamUnavailable(errors.New("nats: stream not found")))
 	assert.False(t, isStreamUnavailable(errors.New("some other error")))
-	assert.False(t, isStreamUnavailable(nats.ErrKeyNotFound))
+	assert.False(t, isStreamUnavailable(jetstream.ErrKeyNotFound))
 }
 
 // --- Terminated instance KV tests ---
@@ -1376,7 +1379,7 @@ func TestJetStreamManager_UpdateTerminatedInstance_NotFound(t *testing.T) {
 	require.NoError(t, jsm.InitTerminatedInstanceBucket())
 
 	_, err = jsm.UpdateTerminatedInstance("i-does-not-exist", func(v *vm.VM) {})
-	assert.ErrorIs(t, err, nats.ErrKeyNotFound)
+	assert.ErrorIs(t, err, jetstream.ErrKeyNotFound)
 }
 
 // TestJetStreamManager_WriteStoppedInstance_OverwritesConcurrentValue verifies
@@ -1416,9 +1419,9 @@ func TestJetStreamManager_WriteStoppedInstance_OverwritesConcurrentValue(t *test
 // deleteTerminatedInstanceBucket deletes the underlying JetStream stream for the terminated-instances KV bucket.
 func deleteTerminatedInstanceBucket(t *testing.T, nc *nats.Conn) {
 	t.Helper()
-	js, err := nc.JetStream()
+	js, err := jetstream.New(nc)
 	require.NoError(t, err)
-	err = js.DeleteStream("KV_" + TerminatedInstanceBucket)
+	err = js.DeleteStream(t.Context(), "KV_"+TerminatedInstanceBucket)
 	require.NoError(t, err)
 }
 
@@ -1504,15 +1507,15 @@ func TestJetStreamManager_InitBuckets_WritesVersion(t *testing.T) {
 	require.NoError(t, jsm.InitClusterStateBucket())
 	require.NoError(t, jsm.InitTerminatedInstanceBucket())
 
-	v, err := utils.ReadVersion(jsm.kv)
+	v, err := kvutil.ReadVersion(t.Context(), jsm.kv)
 	require.NoError(t, err)
 	assert.Equal(t, InstanceStateBucketVersion, v)
 
-	v, err = utils.ReadVersion(jsm.clusterKV)
+	v, err = kvutil.ReadVersion(t.Context(), jsm.clusterKV)
 	require.NoError(t, err)
 	assert.Equal(t, ClusterStateBucketVersion, v)
 
-	v, err = utils.ReadVersion(jsm.terminatedKV)
+	v, err = kvutil.ReadVersion(t.Context(), jsm.terminatedKV)
 	require.NoError(t, err)
 	assert.Equal(t, TerminatedInstanceBucketVersion, v)
 }
@@ -1653,7 +1656,7 @@ func TestJetStreamManager_UpdateMgmtIPAM_ConflictRetry(t *testing.T) {
 	require.NoError(t, jsm.InitClusterStateBucket())
 
 	const subnet = "10.99.8.0/24"
-	defer func() { _ = jsm.clusterKV.Delete(mgmtIPAMKeyPrefix + subnet) }()
+	defer func() { _ = jsm.clusterKV.Delete(t.Context(), mgmtIPAMKeyPrefix+subnet) }()
 
 	_, err = jsm.UpdateMgmtIPAM(subnet, func(r *MgmtIPRecord) {
 		r.Subnet = subnet
@@ -1690,7 +1693,7 @@ func TestJetStreamManager_UpdateMgmtIPAM_ConflictRetry(t *testing.T) {
 }
 
 // TestJetStreamManager_UpdateMgmtIPAM_NotFound verifies createIfAbsent=false
-// surfaces nats.ErrKeyNotFound instead of silently creating a record —
+// surfaces jetstream.ErrKeyNotFound instead of silently creating a record —
 // MgmtIPAllocator.Release relies on this to no-op cleanly when there is
 // nothing to release.
 func TestJetStreamManager_UpdateMgmtIPAM_NotFound(t *testing.T) {
@@ -1703,7 +1706,7 @@ func TestJetStreamManager_UpdateMgmtIPAM_NotFound(t *testing.T) {
 	require.NoError(t, jsm.InitClusterStateBucket())
 
 	_, err = jsm.UpdateMgmtIPAM("10.98.8.0/24-nonexistent", func(r *MgmtIPRecord) {}, false)
-	assert.ErrorIs(t, err, nats.ErrKeyNotFound)
+	assert.ErrorIs(t, err, jetstream.ErrKeyNotFound)
 }
 
 // TestJetStreamManager_UpdateMgmtIPAM_ClusterKVNotInitialized tests error

spinifex/gateway/bedrock/credentials.go

@@ -7,7 +7,8 @@ import (
 	"sync"
 
 	handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
-	"github.com/nats-io/nats.go"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // bedrockCredentialsBucket is the cluster-replicated KV bucket holding
@@ -26,13 +27,13 @@ func credentialKey(accountID, vendor string) string {
 // bedrock-credentials JetStream KV bucket, falling back to a platform-wide
 // default when an account has none. Keys are encrypted at rest.
 type CredentialStore struct {
-	js               nats.JetStreamContext
+	js               jetstream.JetStream
 	masterKey        []byte
 	replicas         int
 	platformDefaults map[string]string
 
 	mu sync.Mutex
-	kv nats.KeyValue
+	kv jetstream.KeyValue
 }
 
 var _ CredentialResolver = (*CredentialStore)(nil)
@@ -40,7 +41,7 @@ var _ CredentialResolver = (*CredentialStore)(nil)
 // NewCredentialStore constructs a CredentialStore. platformDefaults maps
 // vendor to a platform-wide API key used when an account has no key of its
 // own; pass an empty map to disable platform defaults entirely.
-func NewCredentialStore(js nats.JetStreamContext, masterKey []byte, replicas int, platformDefaults map[string]string) *CredentialStore {
+func NewCredentialStore(js jetstream.JetStream, masterKey []byte, replicas int, platformDefaults map[string]string) *CredentialStore {
 	return &CredentialStore{
 		js:               js,
 		masterKey:        masterKey,
@@ -49,23 +50,17 @@ func NewCredentialStore(js nats.JetStreamContext, masterKey []byte, replicas int
 	}
 }
 
-// bucket lazily opens (or creates) the bedrock-credentials KV bucket,
-// caching the handle, mirroring gateway_ecrauth.openSigningBucket.
-func (s *CredentialStore) bucket() (nats.KeyValue, error) {
+// bucket lazily opens (or creates) the cluster-replicated bedrock-credentials
+// KV bucket, caching the handle for subsequent calls.
+func (s *CredentialStore) bucket(ctx context.Context) (jetstream.KeyValue, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if s.kv != nil {
 		return s.kv, nil
 	}
-	kv, err := s.js.CreateKeyValue(&nats.KeyValueConfig{
-		Bucket:   bedrockCredentialsBucket,
-		History:  bedrockCredentialsHistory,
-		Replicas: s.replicas,
-	})
+	kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, s.js, bedrockCredentialsBucket, bedrockCredentialsHistory, s.replicas)
 	if err != nil {
-		if kv, err = s.js.KeyValue(bedrockCredentialsBucket); err != nil {
-			return nil, fmt.Errorf("open %s bucket: %w", bedrockCredentialsBucket, err)
-		}
+		return nil, err
 	}
 	s.kv = kv
 	return kv, nil
@@ -75,13 +70,13 @@ func (s *CredentialStore) bucket() (nats.KeyValue, error) {
 // stored, else the platform default, else ("", false, nil). A nil js is
 // tolerated so a store built with only platformDefaults (e.g. in tests) can
 // still serve the default-only path without touching JetStream.
-func (s *CredentialStore) Resolve(_ context.Context, accountID, vendor string) (string, bool, error) {
+func (s *CredentialStore) Resolve(ctx context.Context, accountID, vendor string) (string, bool, error) {
 	if s.js != nil {
-		kv, err := s.bucket()
+		kv, err := s.bucket(ctx)
 		if err != nil {
 			return "", false, err
 		}
-		entry, err := kv.Get(credentialKey(accountID, vendor))
+		entry, err := kv.Get(ctx, credentialKey(accountID, vendor))
 		switch {
 		case err == nil:
 			plaintext, err := handlers_iam.DecryptSecret(string(entry.Value()), s.masterKey)
@@ -89,7 +84,7 @@ func (s *CredentialStore) Resolve(_ context.Context, accountID, vendor string) (
 				return "", false, fmt.Errorf("decrypt credential for %s/%s: %w", accountID, vendor, err)
 			}
 			return plaintext, true, nil
-		case !errors.Is(err, nats.ErrKeyNotFound):
+		case !errors.Is(err, jetstream.ErrKeyNotFound):
 			return "", false, fmt.Errorf("kv get credential for %s/%s: %w", accountID, vendor, err)
 		}
 	}
@@ -100,8 +95,8 @@ func (s *CredentialStore) Resolve(_ context.Context, accountID, vendor string) (
 }
 
 // PutCredential encrypts and stores key as accountID's vendor credential.
-func (s *CredentialStore) PutCredential(_ context.Context, accountID, vendor, key string) error {
-	kv, err := s.bucket()
+func (s *CredentialStore) PutCredential(ctx context.Context, accountID, vendor, key string) error {
+	kv, err := s.bucket(ctx)
 	if err != nil {
 		return err
 	}
@@ -109,7 +104,7 @@ func (s *CredentialStore) PutCredential(_ context.Context, accountID, vendor, ke
 	if err != nil {
 		return fmt.Errorf("encrypt credential for %s/%s: %w", accountID, vendor, err)
 	}
-	if _, err := kv.Put(credentialKey(accountID, vendor), []byte(ciphertext)); err != nil {
+	if _, err := kv.Put(ctx, credentialKey(accountID, vendor), []byte(ciphertext)); err != nil {
 		return fmt.Errorf("kv put credential for %s/%s: %w", accountID, vendor, err)
 	}
 	return nil

spinifex/gateway/bedrock/credentials_test.go

@@ -17,7 +17,7 @@ func fixedMasterKey() []byte {
 }
 
 // TestCredentialStore_Resolve_NilJS covers the platform-default and
-// vendor-miss paths of a store built with a nil JetStream context (as a
+// vendor-miss paths of a store built with a nil JetStream handle (as a
 // live NATS server is out of scope here). Resolve must consult
 // platformDefaults without dereferencing js.
 func TestCredentialStore_Resolve_NilJS(t *testing.T) {
@@ -62,8 +62,8 @@ func TestNoopCredentialResolver_ResolvesNothing(t *testing.T) {
 // bucket (lazy create), credentialKey, PutCredential (encrypt+Put), and the
 // KV-hit branch of Resolve (Get+decrypt).
 func TestCredentialStore_PutAndResolve_KV(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	store := NewCredentialStore(js, fixedMasterKey(), 1, nil)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	store := NewCredentialStore(testutil.NewJetStream(t, nc), fixedMasterKey(), 1, nil)
 
 	ctx := context.Background()
 	require.NoError(t, store.PutCredential(ctx, "000000000001", "anthropic", "sk-test"))
@@ -75,10 +75,10 @@ func TestCredentialStore_PutAndResolve_KV(t *testing.T) {
 }
 
 // TestCredentialStore_Resolve_KVMissFallsBackToPlatformDefault covers a KV
-// miss (nats.ErrKeyNotFound) that falls through to the platform default.
+// miss (jetstream.ErrKeyNotFound) that falls through to the platform default.
 func TestCredentialStore_Resolve_KVMissFallsBackToPlatformDefault(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	store := NewCredentialStore(js, fixedMasterKey(), 1, map[string]string{"anthropic": "sk-platform-default"})
+	_, nc, _ := testutil.StartTestJetStream(t)
+	store := NewCredentialStore(testutil.NewJetStream(t, nc), fixedMasterKey(), 1, map[string]string{"anthropic": "sk-platform-default"})
 
 	key, ok, err := store.Resolve(context.Background(), "999999999999", "anthropic")
 	require.NoError(t, err)
@@ -89,8 +89,8 @@ func TestCredentialStore_Resolve_KVMissFallsBackToPlatformDefault(t *testing.T)
 // TestCredentialStore_Resolve_KVMissNoPlatformDefault covers a KV miss on an
 // unknown account with no platform defaults configured: ("", false, nil).
 func TestCredentialStore_Resolve_KVMissNoPlatformDefault(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	store := NewCredentialStore(js, fixedMasterKey(), 1, nil)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	store := NewCredentialStore(testutil.NewJetStream(t, nc), fixedMasterKey(), 1, nil)
 
 	key, ok, err := store.Resolve(context.Background(), "999999999999", "anthropic")
 	require.NoError(t, err)

spinifex/gateway/ec2.go

@@ -110,7 +110,7 @@ var ec2Actions = map[string]EC2Handler{
 			return gw.checkPolicyResource(r, "iam", "PassRole", roleARN)
 		}
 		launchQuotaCheck := func() error {
-			return gw.Quota.EnforceLaunch(accountID, aws.StringValue(input.InstanceType), int(aws.Int64Value(input.MaxCount)))
+			return gw.Quota.EnforceLaunch(ctx, accountID, aws.StringValue(input.InstanceType), int(aws.Int64Value(input.MaxCount)))
 		}
 		reservation, err := gateway_ec2_instance.RunInstances(ctx, input, gw.NATSConn, gw.IAMService, accountID, passRoleCheck, launchQuotaCheck, gw.ExpectedNodes)
 		if err != nil {
@@ -118,7 +118,7 @@ var ec2Actions = map[string]EC2Handler{
 		}
 		// Charge the actual launched vCPUs; a counter write failure is drift for
 		// reconcile to correct, so it must not fail the already-successful launch.
-		if err := gw.Quota.ChargeLaunch(accountID, &reservation); err != nil {
+		if err := gw.Quota.ChargeLaunch(ctx, accountID, &reservation); err != nil {
 			slog.WarnContext(ctx, "RunInstances: vcpu quota charge failed, reconcile will correct", "account", accountID, "err", err)
 		}
 		return reservation, nil
@@ -166,7 +166,7 @@ var ec2Actions = map[string]EC2Handler{
 		var delta int
 		if input.InstanceType != nil {
 			resolve := handlers_quota.NATSInstanceTypeResolver(gw.NATSConn, func() int { return gw.ExpectedNodes })
-			d, err := gw.Quota.EnforceRetype(resolve, accountID, aws.StringValue(input.InstanceId), aws.StringValue(input.InstanceType.Value))
+			d, err := gw.Quota.EnforceRetype(ctx, resolve, accountID, aws.StringValue(input.InstanceId), aws.StringValue(input.InstanceType.Value))
 			if err != nil {
 				return nil, err
 			}
@@ -178,7 +178,7 @@ var ec2Actions = map[string]EC2Handler{
 		}
 		// Charge the retype's vCPU growth; a counter write failure is drift for
 		// reconcile to correct, so it must not fail the applied retype.
-		if err := gw.Quota.AddVCPU(accountID, delta); err != nil {
+		if err := gw.Quota.AddVCPU(ctx, accountID, delta); err != nil {
 			slog.WarnContext(ctx, "ModifyInstanceAttribute: vcpu quota charge failed, reconcile will correct", "account", accountID, "err", err)
 		}
 		return out, nil

spinifex/gateway/ec2/instance/RunInstances.go

@@ -101,7 +101,7 @@ func RunInstances(ctx context.Context, input *ec2.RunInstancesInput, natsConn *n
 	}
 	// Hash before any mutation so the same token+params always matches.
 	paramHash := clientTokenParamHash(input)
-	return runInstancesWithClientToken(store, accountID, token, paramHash, func() (ec2.Reservation, error) {
+	return runInstancesWithClientToken(ctx, store, accountID, token, paramHash, func() (ec2.Reservation, error) {
 		return runInstancesInner(ctx, input, natsConn, iamSvc, accountID, passRoleCheck, launchQuotaCheck, expectedNodes)
 	})
 }

spinifex/gateway/ec2/instance/clienttoken.go

@@ -14,6 +14,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 const (
@@ -54,7 +55,7 @@ type tokenRecord struct {
 // ClientTokenStore implements RunInstances ClientToken idempotency over a TTL KV
 // bucket. The first caller owns the launch; duplicates replay (done) or poll (in-flight).
 type ClientTokenStore struct {
-	kv nats.KeyValue
+	kv jetstream.KeyValue
 }
 
 var (
@@ -66,20 +67,24 @@ var (
 // getClientTokenStore lazily initialises the process-wide client-token store via sync.Once.
 func getClientTokenStore(ctx context.Context, nc *nats.Conn) (*ClientTokenStore, error) {
 	ctOnce.Do(func() {
-		js, err := nc.JetStream()
+		js, err := jetstream.New(nc)
 		if err != nil {
 			errCTStoreInit = fmt.Errorf("clienttoken jetstream: %w", err)
 			return
 		}
-		ctStore, errCTStoreInit = newClientTokenStore(js)
+		// The bind happens once per process, so it must not inherit the first
+		// caller's cancellation: a client that disconnects mid-open would poison
+		// the store for every later launch. Deadline-free, so the open falls back
+		// to the JetStream API's own timeout.
+		ctStore, errCTStoreInit = newClientTokenStore(context.WithoutCancel(ctx), js)
 	})
 	return ctStore, errCTStoreInit
 }
 
-func newClientTokenStore(js nats.JetStreamContext) (*ClientTokenStore, error) {
-	kv, err := js.KeyValue(kvBucketClientTokens)
-	if errors.Is(err, nats.ErrBucketNotFound) {
-		kv, err = js.CreateKeyValue(&nats.KeyValueConfig{
+func newClientTokenStore(ctx context.Context, js jetstream.JetStream) (*ClientTokenStore, error) {
+	kv, err := js.KeyValue(ctx, kvBucketClientTokens)
+	if errors.Is(err, jetstream.ErrBucketNotFound) {
+		kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
 			Bucket:  kvBucketClientTokens,
 			History: 1,
 			TTL:     clientTokenTTL,
@@ -112,7 +117,7 @@ func clientTokenParamHash(input *ec2.RunInstancesInput) string {
 // Claim attempts to own the launch for this token.
 // Returns (nil, true, nil) when the caller owns the launch (must Finalize or Abort),
 // (reservation, false, nil) to replay a completed launch, or an error on mismatch/timeout.
-func (c *ClientTokenStore) Claim(accountID, token, paramHash string) (*ec2.Reservation, bool, error) {
+func (c *ClientTokenStore) Claim(ctx context.Context, accountID, token, paramHash string) (*ec2.Reservation, bool, error) {
 	key := clientTokenKey(accountID, token)
 	inflight, err := json.Marshal(tokenRecord{
 		Status:    tokenStatusInFlight,
@@ -122,21 +127,21 @@ func (c *ClientTokenStore) Claim(accountID, token, paramHash string) (*ec2.Reser
 	if err != nil {
 		return nil, false, fmt.Errorf("clienttoken marshal: %w", err)
 	}
-	if _, cerr := c.kv.Create(key, inflight); cerr == nil {
+	if _, cerr := c.kv.Create(ctx, key, inflight); cerr == nil {
 		return nil, true, nil
-	} else if !errors.Is(cerr, nats.ErrKeyExists) {
+	} else if !errors.Is(cerr, jetstream.ErrKeyExists) {
 		return nil, false, fmt.Errorf("clienttoken create %s: %w", key, cerr)
 	}
 
 	deadline := time.Now().Add(clientTokenWaitTimeout)
 	for {
-		entry, gerr := c.kv.Get(key)
+		entry, gerr := c.kv.Get(ctx, key)
 		if gerr != nil {
-			if errors.Is(gerr, nats.ErrKeyNotFound) {
+			if errors.Is(gerr, jetstream.ErrKeyNotFound) {
 				// Owner aborted or record aged out: race for ownership.
-				if _, rcerr := c.kv.Create(key, inflight); rcerr == nil {
+				if _, rcerr := c.kv.Create(ctx, key, inflight); rcerr == nil {
 					return nil, true, nil
-				} else if !errors.Is(rcerr, nats.ErrKeyExists) {
+				} else if !errors.Is(rcerr, jetstream.ErrKeyExists) {
 					return nil, false, fmt.Errorf("clienttoken recreate %s: %w", key, rcerr)
 				}
 				continue
@@ -161,7 +166,7 @@ func (c *ClientTokenStore) Claim(accountID, token, paramHash string) (*ec2.Reser
 }
 
 // Finalize marks the token done with the reservation so duplicates can replay it.
-func (c *ClientTokenStore) Finalize(accountID, token, paramHash string, res *ec2.Reservation) error {
+func (c *ClientTokenStore) Finalize(ctx context.Context, accountID, token, paramHash string, res *ec2.Reservation) error {
 	key := clientTokenKey(accountID, token)
 	data, err := json.Marshal(tokenRecord{
 		Status:      tokenStatusDone,
@@ -172,7 +177,7 @@ func (c *ClientTokenStore) Finalize(accountID, token, paramHash string, res *ec2
 	if err != nil {
 		return fmt.Errorf("clienttoken finalize marshal: %w", err)
 	}
-	if _, err := c.kv.Put(key, data); err != nil {
+	if _, err := c.kv.Put(ctx, key, data); err != nil {
 		return fmt.Errorf("clienttoken finalize put %s: %w", key, err)
 	}
 	return nil
@@ -182,12 +187,13 @@ func (c *ClientTokenStore) Finalize(accountID, token, paramHash string, res *ec2
 // claims the token, replays a completed reservation, or (as owner) launches,
 // finalizes, and aborts on failure. Extracted for unit-testability.
 func runInstancesWithClientToken(
+	ctx context.Context,
 	store *ClientTokenStore,
 	accountID, token, paramHash string,
 	launch func() (ec2.Reservation, error),
 ) (ec2.Reservation, error) {
 	var zero ec2.Reservation
-	replay, owned, cerr := store.Claim(accountID, token, paramHash)
+	replay, owned, cerr := store.Claim(ctx, accountID, token, paramHash)
 	if cerr != nil {
 		if errors.Is(cerr, errIdempotentParamMismatch) {
 			return zero, errors.New(awserrors.ErrorIdempotentParameterMismatch)
@@ -203,11 +209,16 @@ func runInstancesWithClientToken(
 	}
 
 	res, rerr := launch()
+
+	// Recording the launch outcome outlives ctx: a caller that went away mid-launch
+	// is exactly when the record must be settled, and leaving it in-flight parks
+	// every retry of that token behind the poll deadline until the record ages out.
+	outcomeCtx := context.WithoutCancel(ctx)
 	if rerr != nil {
-		store.Abort(accountID, token)
+		store.Abort(outcomeCtx, accountID, token)
 		return zero, rerr
 	}
-	if ferr := store.Finalize(accountID, token, paramHash, &res); ferr != nil {
+	if ferr := store.Finalize(outcomeCtx, accountID, token, paramHash, &res); ferr != nil {
 		// Launch succeeded; finalize failure only weakens future dedup — don't fail.
 		slog.Warn("RunInstances: failed to finalize client-token record", "token", token, "err", ferr)
 	}
@@ -215,9 +226,9 @@ func runInstancesWithClientToken(
 }
 
 // Abort drops the in-flight token so a retry can re-launch.
-func (c *ClientTokenStore) Abort(accountID, token string) {
+func (c *ClientTokenStore) Abort(ctx context.Context, accountID, token string) {
 	key := clientTokenKey(accountID, token)
-	if err := c.kv.Delete(key); err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
+	if err := c.kv.Delete(ctx, key); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
 		slog.Warn("clienttoken: failed to abort token record", "key", key, "err", err)
 	}
 }

spinifex/gateway/ec2/instance/clienttoken_test.go

@@ -1,7 +1,6 @@
 package gateway_ec2_instance
 
 import (
-	"context"
 	"errors"
 	"sync"
 	"sync/atomic"
@@ -20,8 +19,8 @@ const ctTestAccount = "111122223333"
 
 func newTestClientTokenStore(t *testing.T) *ClientTokenStore {
 	t.Helper()
-	_, _, js := testutil.StartTestJetStream(t)
-	store, err := newClientTokenStore(js)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	store, err := newClientTokenStore(t.Context(), testutil.NewJetStream(t, nc))
 	require.NoError(t, err)
 	return store
 }
@@ -32,14 +31,14 @@ func TestClientToken_ReplaysCompletedReservation(t *testing.T) {
 	store := newTestClientTokenStore(t)
 	const tok, hash = "tok-1", "hash-a"
 
-	_, owned, err := store.Claim(ctTestAccount, tok, hash)
+	_, owned, err := store.Claim(t.Context(), ctTestAccount, tok, hash)
 	require.NoError(t, err)
 	require.True(t, owned, "first caller owns the launch")
 
 	res := &ec2.Reservation{ReservationId: aws.String("r-123")}
-	require.NoError(t, store.Finalize(ctTestAccount, tok, hash, res))
+	require.NoError(t, store.Finalize(t.Context(), ctTestAccount, tok, hash, res))
 
-	replay, owned2, err := store.Claim(ctTestAccount, tok, hash)
+	replay, owned2, err := store.Claim(t.Context(), ctTestAccount, tok, hash)
 	require.NoError(t, err)
 	assert.False(t, owned2, "duplicate caller does not own the launch")
 	require.NotNil(t, replay)
@@ -51,12 +50,12 @@ func TestClientToken_ParamMismatchRejected(t *testing.T) {
 	store := newTestClientTokenStore(t)
 	const tok = "tok-2"
 
-	_, owned, err := store.Claim(ctTestAccount, tok, "hash-a")
+	_, owned, err := store.Claim(t.Context(), ctTestAccount, tok, "hash-a")
 	require.NoError(t, err)
 	require.True(t, owned)
-	require.NoError(t, store.Finalize(ctTestAccount, tok, "hash-a", &ec2.Reservation{ReservationId: aws.String("r-1")}))
+	require.NoError(t, store.Finalize(t.Context(), ctTestAccount, tok, "hash-a", &ec2.Reservation{ReservationId: aws.String("r-1")}))
 
-	_, _, err = store.Claim(ctTestAccount, tok, "hash-DIFFERENT")
+	_, _, err = store.Claim(t.Context(), ctTestAccount, tok, "hash-DIFFERENT")
 	assert.ErrorIs(t, err, errIdempotentParamMismatch)
 }
 
@@ -66,13 +65,13 @@ func TestClientToken_AbortAllowsRelaunch(t *testing.T) {
 	store := newTestClientTokenStore(t)
 	const tok, hash = "tok-3", "hash-a"
 
-	_, owned, err := store.Claim(ctTestAccount, tok, hash)
+	_, owned, err := store.Claim(t.Context(), ctTestAccount, tok, hash)
 	require.NoError(t, err)
 	require.True(t, owned)
 
-	store.Abort(ctTestAccount, tok)
+	store.Abort(t.Context(), ctTestAccount, tok)
 
-	_, owned2, err := store.Claim(ctTestAccount, tok, hash)
+	_, owned2, err := store.Claim(t.Context(), ctTestAccount, tok, hash)
 	require.NoError(t, err)
 	assert.True(t, owned2, "after abort the token is free to re-own")
 }
@@ -93,7 +92,7 @@ func TestClientToken_ConcurrentSingleOwner(t *testing.T) {
 		wg.Add(1)
 		go func(i int) {
 			defer wg.Done()
-			replay, owned, err := store.Claim(ctTestAccount, tok, hash)
+			replay, owned, err := store.Claim(t.Context(), ctTestAccount, tok, hash)
 			if err != nil {
 				errs[i] = err
 				return
@@ -101,7 +100,7 @@ func TestClientToken_ConcurrentSingleOwner(t *testing.T) {
 			if owned {
 				atomic.AddInt32(&owners, 1)
 				// Simulate the launch then publish the reservation.
-				errs[i] = store.Finalize(ctTestAccount, tok, hash, res)
+				errs[i] = store.Finalize(t.Context(), ctTestAccount, tok, hash, res)
 				replays[i] = res
 				return
 			}
@@ -147,13 +146,13 @@ func TestClientTokenParamHash_IgnoresTokenReflectsParams(t *testing.T) {
 func TestRunInstancesWithClientToken_ReplaySkipsLaunch(t *testing.T) {
 	store := newTestClientTokenStore(t)
 	const tok, hash = "rt-1", "h"
-	_, owned, err := store.Claim(ctTestAccount, tok, hash)
+	_, owned, err := store.Claim(t.Context(), ctTestAccount, tok, hash)
 	require.NoError(t, err)
 	require.True(t, owned)
-	require.NoError(t, store.Finalize(ctTestAccount, tok, hash, &ec2.Reservation{ReservationId: aws.String("r-x")}))
+	require.NoError(t, store.Finalize(t.Context(), ctTestAccount, tok, hash, &ec2.Reservation{ReservationId: aws.String("r-x")}))
 
 	launched := false
-	res, err := runInstancesWithClientToken(store, ctTestAccount, tok, hash, func() (ec2.Reservation, error) {
+	res, err := runInstancesWithClientToken(t.Context(), store, ctTestAccount, tok, hash, func() (ec2.Reservation, error) {
 		launched = true
 		return ec2.Reservation{}, nil
 	})
@@ -172,11 +171,11 @@ func TestRunInstancesWithClientToken_OwnerLaunchesOnceThenReplay(t *testing.T) {
 		return ec2.Reservation{ReservationId: aws.String("r-own")}, nil
 	}
 
-	res, err := runInstancesWithClientToken(store, ctTestAccount, tok, hash, launch)
+	res, err := runInstancesWithClientToken(t.Context(), store, ctTestAccount, tok, hash, launch)
 	require.NoError(t, err)
 	assert.Equal(t, "r-own", aws.StringValue(res.ReservationId))
 
-	res2, err := runInstancesWithClientToken(store, ctTestAccount, tok, hash, launch)
+	res2, err := runInstancesWithClientToken(t.Context(), store, ctTestAccount, tok, hash, launch)
 	require.NoError(t, err)
 	assert.Equal(t, "r-own", aws.StringValue(res2.ReservationId))
 	assert.Equal(t, 1, launches, "duplicate must replay, not relaunch")
@@ -187,13 +186,13 @@ func TestRunInstancesWithClientToken_LaunchFailureAborts(t *testing.T) {
 	store := newTestClientTokenStore(t)
 	const tok, hash = "rt-3", "h"
 
-	_, err := runInstancesWithClientToken(store, ctTestAccount, tok, hash, func() (ec2.Reservation, error) {
+	_, err := runInstancesWithClientToken(t.Context(), store, ctTestAccount, tok, hash, func() (ec2.Reservation, error) {
 		return ec2.Reservation{}, errors.New("no capacity")
 	})
 	require.Error(t, err)
 
 	relaunched := false
-	res, err := runInstancesWithClientToken(store, ctTestAccount, tok, hash, func() (ec2.Reservation, error) {
+	res, err := runInstancesWithClientToken(t.Context(), store, ctTestAccount, tok, hash, func() (ec2.Reservation, error) {
 		relaunched = true
 		return ec2.Reservation{ReservationId: aws.String("r-retry")}, nil
 	})
@@ -207,13 +206,13 @@ func TestRunInstancesWithClientToken_LaunchFailureAborts(t *testing.T) {
 func TestRunInstancesWithClientToken_ParamMismatchMapsAWSError(t *testing.T) {
 	store := newTestClientTokenStore(t)
 	const tok = "rt-4"
-	_, owned, err := store.Claim(ctTestAccount, tok, "hA")
+	_, owned, err := store.Claim(t.Context(), ctTestAccount, tok, "hA")
 	require.NoError(t, err)
 	require.True(t, owned)
-	require.NoError(t, store.Finalize(ctTestAccount, tok, "hA", &ec2.Reservation{ReservationId: aws.String("r")}))
+	require.NoError(t, store.Finalize(t.Context(), ctTestAccount, tok, "hA", &ec2.Reservation{ReservationId: aws.String("r")}))
 
 	launched := false
-	_, err = runInstancesWithClientToken(store, ctTestAccount, tok, "hB", func() (ec2.Reservation, error) {
+	_, err = runInstancesWithClientToken(t.Context(), store, ctTestAccount, tok, "hB", func() (ec2.Reservation, error) {
 		launched = true
 		return ec2.Reservation{}, nil
 	})
@@ -225,10 +224,10 @@ func TestRunInstancesWithClientToken_ParamMismatchMapsAWSError(t *testing.T) {
 // instance on subsequent calls.
 func TestGetClientTokenStore_BindsOnce(t *testing.T) {
 	_, nc, _ := testutil.StartTestJetStream(t)
-	s1, err := getClientTokenStore(context.Background(), nc)
+	s1, err := getClientTokenStore(t.Context(), nc)
 	require.NoError(t, err)
 	require.NotNil(t, s1)
-	s2, err := getClientTokenStore(context.Background(), nc)
+	s2, err := getClientTokenStore(t.Context(), nc)
 	require.NoError(t, err)
 	assert.Same(t, s1, s2, "store binds once")
 }
@@ -239,7 +238,7 @@ func TestClientToken_InFlightWaitTimesOut(t *testing.T) {
 	store := newTestClientTokenStore(t)
 	const tok, hash = "wt-1", "h"
 
-	_, owned, err := store.Claim(ctTestAccount, tok, hash)
+	_, owned, err := store.Claim(t.Context(), ctTestAccount, tok, hash)
 	require.NoError(t, err)
 	require.True(t, owned, "owner holds the in-flight record and never finalizes")
 
@@ -247,6 +246,6 @@ func TestClientToken_InFlightWaitTimesOut(t *testing.T) {
 	clientTokenWaitTimeout, clientTokenPollStep = 30*time.Millisecond, 10*time.Millisecond
 	defer func() { clientTokenWaitTimeout, clientTokenPollStep = origTimeout, origStep }()
 
-	_, _, err = store.Claim(ctTestAccount, tok, hash)
+	_, _, err = store.Claim(t.Context(), ctTestAccount, tok, hash)
 	assert.ErrorIs(t, err, errClientTokenWaitTimeout)
 }

spinifex/gateway/ec2/spotinstance/RequestSpotInstances.go

@@ -76,7 +76,7 @@ func RequestSpotInstances(ctx context.Context, input *ec2.RequestSpotInstancesIn
 	// Gate on the per-account vCPU cap exactly like the on-demand path; spot
 	// launches are real VMs and must not slip past the quota.
 	launchQuotaCheck := func() error {
-		return quota.EnforceLaunch(accountID, aws.StringValue(runInput.InstanceType), int(count))
+		return quota.EnforceLaunch(ctx, accountID, aws.StringValue(runInput.InstanceType), int(count))
 	}
 
 	// RunInstances normalises runInput in place (e.g. instance profile to ARN),
@@ -88,7 +88,7 @@ func RequestSpotInstances(ctx context.Context, input *ec2.RequestSpotInstancesIn
 
 	// Charge the actual launched vCPUs; a counter write failure is drift for
 	// reconcile to correct, so it must not fail the already-successful launch.
-	if err := quota.ChargeLaunch(accountID, &reservation); err != nil {
+	if err := quota.ChargeLaunch(ctx, accountID, &reservation); err != nil {
 		slog.WarnContext(ctx, "RequestSpotInstances: vcpu quota charge failed, reconcile will correct", "account", accountID, "err", err)
 	}
 

spinifex/gateway/ecr/registry_image_test.go

@@ -157,7 +157,8 @@ func TestListImages_TaggedAndUntagged(t *testing.T) {
 // canonicalize on meta.Digest to emit a real digest and keep tag association.
 // The memory store returns real digests and would mask this.
 func TestListImages_KVDigestForm(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 	reg := NewRegistry(objectstore.NewMemoryObjectStore(), ecr.NewKVMetaStore(js), testAccount)
 	repo := "team/app"
 	require.NoError(t, reg.Meta.PutRepo(context.Background(), testAccount, ecr.RepoMeta{Name: repo}))

spinifex/gateway/ecr_createdeleterepository_test.go

@@ -20,7 +20,8 @@ import (
 // subjects against an embedded KV-backed service, returning a gateway client.
 func newRepoLifecycleGateway(t *testing.T) (*GatewayConfig, *nats.Conn) {
 	t.Helper()
-	_, nc, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 	svc := handlers_ecr.NewKVMetaService(js)
 	serveECRMeta(t, nc, handlers_ecr.SubjectRepoCreate, svc.RepoCreate)
 	serveECRMeta(t, nc, handlers_ecr.SubjectRepoDescribe, svc.RepoDescribe)

spinifex/gateway/ecr_describerepositories_test.go

@@ -42,7 +42,8 @@ func serveECRMeta[I any, O any](t *testing.T, nc *nats.Conn, subject string, fn
 
 func newDescribeReposGateway(t *testing.T, repos ...string) *GatewayConfig {
 	t.Helper()
-	_, nc, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 	svc := handlers_ecr.NewKVMetaService(js)
 	serveECRMeta(t, nc, handlers_ecr.SubjectRepoCreate, svc.RepoCreate)
 	serveECRMeta(t, nc, handlers_ecr.SubjectRepoDescribe, svc.RepoDescribe)

spinifex/gateway/ecrapi/lifecyclepolicy_test.go

@@ -20,7 +20,8 @@ const validLifecyclePolicy = `{"rules":[{"rulePriority":1,"selection":{"tagStatu
 // lifecycle metadata subjects served, returning the gateway NATS connection.
 func newLifecycleTestConn(t *testing.T) *nats.Conn {
 	t.Helper()
-	_, nc, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 	svc := handlers_ecr.NewKVMetaService(js)
 	serveMeta(t, nc, handlers_ecr.SubjectRepoCreate, svc.RepoCreate)
 	serveMeta(t, nc, handlers_ecr.SubjectRepoDescribe, svc.RepoDescribe)

spinifex/gateway/ecrapi/repositorypolicy_test.go

@@ -45,7 +45,8 @@ func serveMeta[I any, O any](t *testing.T, nc *nats.Conn, subject string, fn fun
 // metadata subjects served, returning the gateway NATS connection.
 func newPolicyTestConn(t *testing.T) *nats.Conn {
 	t.Helper()
-	_, nc, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 	svc := handlers_ecr.NewKVMetaService(js)
 	serveMeta(t, nc, handlers_ecr.SubjectRepoCreate, svc.RepoCreate)
 	serveMeta(t, nc, handlers_ecr.SubjectRepoDescribe, svc.RepoDescribe)

spinifex/gateway/ecrauth/ecrauth_test.go

@@ -37,42 +37,46 @@ func samplePrincipal() Principal {
 }
 
 func TestLoadOrCreateSigningKey_CreatesThenReloads(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 
-	key1, verify1, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	key1, verify1, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 	require.NotEmpty(t, key1.Kid)
 	require.Contains(t, verify1, key1.Kid)
 
 	// Second call must reload the same persisted key, not mint a new one.
-	key2, verify2, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	key2, verify2, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 	assert.Equal(t, key1.Kid, key2.Kid, "persisted signing key must be reused")
 	assert.Len(t, verify2, 1)
 }
 
 func TestLoadOrCreateSigningKey_StoresEncrypted(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	key, _, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	key, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 
-	kv, err := js.KeyValue(SigningBucket)
+	kv, err := js.KeyValue(t.Context(), SigningBucket)
 	require.NoError(t, err)
-	entry, err := kv.Get(signingKeyName(key.Kid))
+	entry, err := kv.Get(t.Context(), signingKeyName(key.Kid))
 	require.NoError(t, err)
 	assert.NotContains(t, string(entry.Value()), "PRIVATE KEY",
 		"signing key must be encrypted at rest")
 }
 
 func TestLoadOrCreateSigningKey_EmptyMasterKeyRejected(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	_, _, err := LoadOrCreateSigningKey(js, nil, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	_, _, err := LoadOrCreateSigningKey(t.Context(), js, nil, 1)
 	require.Error(t, err)
 }
 
 func TestIssuerVerifier_RoundTrip(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	key, verify, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 
 	iss := NewIssuer(key, testAudience)
@@ -88,8 +92,9 @@ func TestIssuerVerifier_RoundTrip(t *testing.T) {
 }
 
 func TestIssuer_MintRequiresAccountID(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	key, _, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	key, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 
 	p := samplePrincipal()
@@ -103,8 +108,9 @@ func TestIssuer_MintRequiresAccountID(t *testing.T) {
 // accessKeyID, an unsupported principalType) must fail to mint at all, since
 // a token that can't name a lookup key can never be safely rehydrated.
 func TestMint_RequiresCompleteIdentityPointer(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	key, _, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	key, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 	iss := NewIssuer(key, testAudience)
 
@@ -136,8 +142,9 @@ func TestSupportedPrincipalType(t *testing.T) {
 }
 
 func TestVerifier_RejectsWrongAudience(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	key, verify, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 
 	tok, _, err := NewIssuer(key, testAudience).Mint(samplePrincipal())
@@ -148,8 +155,9 @@ func TestVerifier_RejectsWrongAudience(t *testing.T) {
 }
 
 func TestVerifier_RejectsUnknownKid(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	key, _, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	key, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 
 	tok, _, err := NewIssuer(key, testAudience).Mint(samplePrincipal())
@@ -161,8 +169,9 @@ func TestVerifier_RejectsUnknownKid(t *testing.T) {
 }
 
 func TestVerifier_RejectsExpiredToken(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	key, verify, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 
 	// Hand-mint an already-expired token with the same signing key.
@@ -188,8 +197,9 @@ func TestVerifier_RejectsExpiredToken(t *testing.T) {
 // Mint's own checks, proving Verify itself refuses a claim set missing any
 // part of the identity pointer even if it were signed by some other path.
 func TestVerifier_RejectsIncompleteIdentityPointer(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	key, verify, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 
 	base := func() Claims {
@@ -234,8 +244,9 @@ func TestVerifier_RejectsIncompleteIdentityPointer(t *testing.T) {
 }
 
 func TestVerifier_RejectsNonES256(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	_, verify, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	_, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 
 	// HS256 token must be refused regardless of signature.

spinifex/gateway/ecrauth/rotator.go

@@ -6,7 +6,7 @@ import (
 	"sort"
 	"time"
 
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Rotation defaults (ecr-v1 Q3): rotate the active signing key every 30 days,
@@ -17,6 +17,10 @@ const (
 	DefaultRetainFor             = 24 * time.Hour
 	DefaultRetainCount           = 2
 	DefaultRotationCheckInterval = time.Hour
+
+	// rotationCycleTimeout bounds the KV work of one rotation cycle, so a wedged
+	// JetStream stalls a single tick instead of the whole loop until shutdown.
+	rotationCycleTimeout = 30 * time.Second
 )
 
 // rotationParams holds the rotation/retention policy.
@@ -61,7 +65,7 @@ func planRotation(metas []keyMeta, now time.Time, p rotationParams) (mint bool,
 // safe to run on every awsgw instance — mint and delete on the shared KV bucket
 // converge, and a delete of an already-gone key is a no-op.
 type Rotator struct {
-	kv        nats.KeyValue
+	kv        jetstream.KeyValue
 	masterKey []byte
 	issuer    *Issuer
 	verifier  *Verifier
@@ -73,8 +77,8 @@ type Rotator struct {
 // NewRotator opens the signing-key bucket and builds a rotator that keeps issuer
 // and verifier current. replicas matches the cluster size for first-run bucket
 // creation.
-func NewRotator(js nats.JetStreamContext, masterKey []byte, replicas int, issuer *Issuer, verifier *Verifier) (*Rotator, error) {
-	kv, err := openSigningBucket(js, masterKey, replicas)
+func NewRotator(ctx context.Context, js jetstream.JetStream, masterKey []byte, replicas int, issuer *Issuer, verifier *Verifier) (*Rotator, error) {
+	kv, err := openSigningBucket(ctx, js, masterKey, replicas)
 	if err != nil {
 		return nil, err
 	}
@@ -105,15 +109,17 @@ func (r *Rotator) Run(ctx context.Context) {
 			slog.Info("ECR signing-key rotator stopped")
 			return
 		case <-ticker.C:
-			r.rotateOnce(r.now())
+			cycleCtx, cancel := context.WithTimeout(ctx, rotationCycleTimeout)
+			r.rotateOnce(cycleCtx, r.now())
+			cancel()
 		}
 	}
 }
 
 // rotateOnce runs one rotation/prune cycle and refreshes the issuer/verifier.
 // Errors are logged and skipped so a transient KV fault never crashes the loop.
-func (r *Rotator) rotateOnce(now time.Time) {
-	_, _, metas, err := reloadKeys(r.kv, r.masterKey)
+func (r *Rotator) rotateOnce(ctx context.Context, now time.Time) {
+	_, _, metas, err := reloadKeys(ctx, r.kv, r.masterKey)
 	if err != nil {
 		slog.Warn("ECR signing-key rotation: load keys failed", "err", err)
 		return
@@ -121,7 +127,7 @@ func (r *Rotator) rotateOnce(now time.Time) {
 
 	mint, prune := planRotation(metas, now, r.params)
 	if mint {
-		newKey, err := generateSigningKey(r.kv, r.masterKey)
+		newKey, err := generateSigningKey(ctx, r.kv, r.masterKey)
 		if err != nil {
 			slog.Warn("ECR signing-key rotation: mint failed", "err", err)
 			return
@@ -129,19 +135,19 @@ func (r *Rotator) rotateOnce(now time.Time) {
 		slog.Info("ECR signing-key rotated", "kid", newKey.Kid)
 	}
 	for _, kid := range prune {
-		if err := deleteSigningKey(r.kv, kid); err != nil {
+		if err := deleteSigningKey(ctx, r.kv, kid); err != nil {
 			slog.Warn("ECR signing-key rotation: prune failed", "kid", kid, "err", err)
 			continue
 		}
 		slog.Info("ECR signing-key pruned", "kid", kid)
 	}
 
-	r.refresh()
+	r.refresh(ctx)
 }
 
 // refresh reloads the canonical key set and installs it on the issuer/verifier.
-func (r *Rotator) refresh() {
-	active, verify, _, err := reloadKeys(r.kv, r.masterKey)
+func (r *Rotator) refresh(ctx context.Context) {
+	active, verify, _, err := reloadKeys(ctx, r.kv, r.masterKey)
 	if err != nil {
 		slog.Warn("ECR signing-key rotation: refresh failed", "err", err)
 		return

spinifex/gateway/ecrauth/rotator_test.go

@@ -77,13 +77,14 @@ func TestPlanRotation(t *testing.T) {
 // always due: a new active key is minted, installed on the issuer, and accepted
 // by the verifier, while the prior key stays verifiable.
 func TestRotateOnce_MintsAndSwaps(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	key, verify, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 
 	issuer := NewIssuer(key, testAudience)
 	verifier := NewVerifier(verify, testAudience)
-	kv, err := openSigningBucket(js, testMasterKey, 1)
+	kv, err := openSigningBucket(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 
 	// Token minted under the original key before rotation.
@@ -95,7 +96,7 @@ func TestRotateOnce_MintsAndSwaps(t *testing.T) {
 		params: rotationParams{rotateAfter: 0, retainFor: time.Hour, retainCount: 2},
 		now:    func() time.Time { return time.Now().UTC() },
 	}
-	r.rotateOnce(time.Now().UTC())
+	r.rotateOnce(t.Context(), time.Now().UTC())
 
 	// A new active key took over.
 	assert.NotEqual(t, key.Kid, issuer.ActiveKid(), "issuer should sign with the rotated key")
@@ -108,7 +109,7 @@ func TestRotateOnce_MintsAndSwaps(t *testing.T) {
 	_, err = verifier.Verify(newTok)
 	require.NoError(t, err, "token from the rotated key must verify")
 
-	_, _, metas, err := reloadKeys(kv, testMasterKey)
+	_, _, metas, err := reloadKeys(t.Context(), kv, testMasterKey)
 	require.NoError(t, err)
 	assert.Len(t, metas, 2, "previous key retained alongside the new active key")
 }
@@ -117,12 +118,13 @@ func TestRotateOnce_MintsAndSwaps(t *testing.T) {
 // retainFor=0 (and no mint) so the older key is pruned and dropped from the
 // verifier.
 func TestRotateOnce_PrunesExpiredKey(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	first, verify, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	first, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
-	kv, err := openSigningBucket(js, testMasterKey, 1)
+	kv, err := openSigningBucket(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
-	second, err := generateSigningKey(kv, testMasterKey)
+	second, err := generateSigningKey(t.Context(), kv, testMasterKey)
 	require.NoError(t, err)
 	require.NotEqual(t, first.Kid, second.Kid)
 
@@ -134,9 +136,9 @@ func TestRotateOnce_PrunesExpiredKey(t *testing.T) {
 		params: rotationParams{rotateAfter: 1000 * time.Hour, retainFor: 0, retainCount: 2},
 		now:    func() time.Time { return time.Now().UTC() },
 	}
-	r.rotateOnce(time.Now().UTC())
+	r.rotateOnce(t.Context(), time.Now().UTC())
 
-	active, _, metas, err := reloadKeys(kv, testMasterKey)
+	active, _, metas, err := reloadKeys(t.Context(), kv, testMasterKey)
 	require.NoError(t, err)
 	require.Len(t, metas, 1, "the older key should be pruned")
 	assert.Equal(t, second.Kid, active.Kid, "newest key remains active")
@@ -147,12 +149,13 @@ func TestRotateOnce_PrunesExpiredKey(t *testing.T) {
 // with rotateAfter=0 so it mints on the first tick, then cancels: the active key
 // changes and Run returns.
 func TestRotator_RunRotatesThenStops(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	key, verify, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 	issuer := NewIssuer(key, testAudience)
 	verifier := NewVerifier(verify, testAudience)
-	kv, err := openSigningBucket(js, testMasterKey, 1)
+	kv, err := openSigningBucket(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
 
 	r := &Rotator{
@@ -180,14 +183,15 @@ func TestRotator_RunRotatesThenStops(t *testing.T) {
 // TestIssuerVerifier_ConcurrentSwap exercises the hot-swap under -race: Mint and
 // Verify run while SetActiveKey/SetKeys rotate the key set.
 func TestIssuerVerifier_ConcurrentSwap(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
-	k1, v1, err := LoadOrCreateSigningKey(js, testMasterKey, 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	k1, v1, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
-	kv, err := openSigningBucket(js, testMasterKey, 1)
+	kv, err := openSigningBucket(t.Context(), js, testMasterKey, 1)
 	require.NoError(t, err)
-	k2, err := generateSigningKey(kv, testMasterKey)
+	k2, err := generateSigningKey(t.Context(), kv, testMasterKey)
 	require.NoError(t, err)
-	_, v2, _, err := reloadKeys(kv, testMasterKey)
+	_, v2, _, err := reloadKeys(t.Context(), kv, testMasterKey)
 	require.NoError(t, err)
 
 	issuer := NewIssuer(k1, testAudience)

spinifex/gateway/ecrauth/signingkey.go

@@ -6,6 +6,7 @@
 package gateway_ecrauth
 
 import (
+	"context"
 	"crypto/ecdsa"
 	"crypto/elliptic"
 	"crypto/rand"
@@ -19,7 +20,8 @@ import (
 	"time"
 
 	handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
-	"github.com/nats-io/nats.go"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 const (
@@ -73,17 +75,17 @@ func kidFor(pub *ecdsa.PublicKey) (string, error) {
 // returns the active signing key. On first run it generates and persists one.
 // The active key is the newest by KV creation time (ties broken by kid), so all
 // nodes converge on the same signer and a freshly rotated key takes over.
-func LoadOrCreateSigningKey(js nats.JetStreamContext, masterKey []byte, replicas int) (*SigningKey, map[string]*ecdsa.PublicKey, error) {
-	kv, err := openSigningBucket(js, masterKey, replicas)
+func LoadOrCreateSigningKey(ctx context.Context, js jetstream.JetStream, masterKey []byte, replicas int) (*SigningKey, map[string]*ecdsa.PublicKey, error) {
+	kv, err := openSigningBucket(ctx, js, masterKey, replicas)
 	if err != nil {
 		return nil, nil, err
 	}
-	active, verify, _, err := reloadKeys(kv, masterKey)
+	active, verify, _, err := reloadKeys(ctx, kv, masterKey)
 	if err != nil {
 		return nil, nil, err
 	}
 	if active == nil {
-		active, err = generateSigningKey(kv, masterKey)
+		active, err = generateSigningKey(ctx, kv, masterKey)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -94,32 +96,22 @@ func LoadOrCreateSigningKey(js nats.JetStreamContext, masterKey []byte, replicas
 
 // openSigningBucket validates inputs and returns the awsgw-keys KV handle,
 // creating the cluster-replicated bucket on first use.
-func openSigningBucket(js nats.JetStreamContext, masterKey []byte, replicas int) (nats.KeyValue, error) {
+func openSigningBucket(ctx context.Context, js jetstream.JetStream, masterKey []byte, replicas int) (jetstream.KeyValue, error) {
 	if js == nil {
 		return nil, errors.New("ecrauth: nil JetStream context")
 	}
 	if len(masterKey) == 0 {
 		return nil, errors.New("ecrauth: empty master key")
 	}
-	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{
-		Bucket:   SigningBucket,
-		History:  signingKeyHistory,
-		Replicas: replicas,
-	})
-	if err != nil {
-		if kv, err = js.KeyValue(SigningBucket); err != nil {
-			return nil, fmt.Errorf("open %s bucket: %w", SigningBucket, err)
-		}
-	}
-	return kv, nil
+	return kvutil.GetOrCreateBucketWithReplicas(ctx, js, SigningBucket, signingKeyHistory, replicas)
 }
 
 // reloadKeys reads every stored signing key, returning the verify set, each
 // key's metadata, and the active (newest) decrypted key. active is nil for an
 // empty bucket. Used by both startup and the rotation scheduler.
-func reloadKeys(kv nats.KeyValue, masterKey []byte) (active *SigningKey, verify map[string]*ecdsa.PublicKey, metas []keyMeta, err error) {
-	names, err := kv.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+func reloadKeys(ctx context.Context, kv jetstream.KeyValue, masterKey []byte) (active *SigningKey, verify map[string]*ecdsa.PublicKey, metas []keyMeta, err error) {
+	names, err := kvutil.Keys(ctx, kv)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, nil, nil, fmt.Errorf("list signing keys: %w", err)
 	}
 
@@ -129,7 +121,7 @@ func reloadKeys(kv nats.KeyValue, masterKey []byte) (active *SigningKey, verify
 		if !strings.HasPrefix(name, signingKeyPrefix) {
 			continue
 		}
-		entry, err := kv.Get(name)
+		entry, err := kv.Get(ctx, name)
 		if err != nil {
 			return nil, nil, nil, fmt.Errorf("kv get %s: %w", name, err)
 		}
@@ -149,8 +141,8 @@ func reloadKeys(kv nats.KeyValue, masterKey []byte) (active *SigningKey, verify
 
 // deleteSigningKey removes a rotated-out key from the bucket once its retention
 // window has elapsed.
-func deleteSigningKey(kv nats.KeyValue, kid string) error {
-	if err := kv.Delete(signingKeyName(kid)); err != nil {
+func deleteSigningKey(ctx context.Context, kv jetstream.KeyValue, kid string) error {
+	if err := kv.Delete(ctx, signingKeyName(kid)); err != nil {
 		return fmt.Errorf("kv delete %s: %w", signingKeyName(kid), err)
 	}
 	return nil
@@ -172,7 +164,7 @@ func decodeSigningKey(kid string, ciphertext []byte, masterKey []byte) (*Signing
 
 // generateSigningKey creates a fresh ES256 key, persists the encrypted PEM under
 // jwt-signing/{kid}, and returns it.
-func generateSigningKey(kv nats.KeyValue, masterKey []byte) (*SigningKey, error) {
+func generateSigningKey(ctx context.Context, kv jetstream.KeyValue, masterKey []byte) (*SigningKey, error) {
 	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
 	if err != nil {
 		return nil, fmt.Errorf("generate ES256 key: %w", err)
@@ -190,7 +182,7 @@ func generateSigningKey(kv nats.KeyValue, masterKey []byte) (*SigningKey, error)
 	if err != nil {
 		return nil, fmt.Errorf("encrypt signing key: %w", err)
 	}
-	if _, err := kv.Put(signingKeyName(kid), []byte(ciphertext)); err != nil {
+	if _, err := kv.Put(ctx, signingKeyName(kid), []byte(ciphertext)); err != nil {
 		return nil, fmt.Errorf("kv put %s: %w", signingKeyName(kid), err)
 	}
 	return &SigningKey{Kid: kid, priv: priv}, nil

spinifex/gateway/ecrauth_middleware_test.go

@@ -32,8 +32,9 @@ func ecrTestMasterKey(t *testing.T) []byte {
 // newECRAuth builds a wired issuer+verifier over a freshly generated signing key.
 func newECRAuth(t *testing.T) (*gateway_ecrauth.Issuer, *gateway_ecrauth.Verifier) {
 	t.Helper()
-	_, _, js := testutil.StartTestJetStream(t)
-	key, verify, err := gateway_ecrauth.LoadOrCreateSigningKey(js, ecrTestMasterKey(t), 1)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	key, verify, err := gateway_ecrauth.LoadOrCreateSigningKey(t.Context(), js, ecrTestMasterKey(t), 1)
 	require.NoError(t, err)
 	return gateway_ecrauth.NewIssuer(key, ecrTestAudience), gateway_ecrauth.NewVerifier(verify, ecrTestAudience)
 }

spinifex/gateway/eks/token_review_test.go

@@ -10,6 +10,7 @@ import (
 	handlers_eks "github.com/mulgadc/spinifex/spinifex/handlers/eks"
 	"github.com/mulgadc/spinifex/spinifex/testutil"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -56,9 +57,11 @@ func TestWebhookTokenReview_InfraFaultIsServerInternal(t *testing.T) {
 }
 
 func TestWebhookTokenReview_ResolvesIdentity(t *testing.T) {
-	_, nc, js := testutil.StartTestJetStream(t)
-	kv := testutil.SeedKV(t, js, handlers_eks.AccountBucketName("111122223333"), nil)
-	require.NoError(t, handlers_eks.PutAccessEntryRecord(kv, &handlers_eks.AccessEntryRecord{
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	kv, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{Bucket: handlers_eks.AccountBucketName("111122223333")})
+	require.NoError(t, err)
+	require.NoError(t, handlers_eks.PutAccessEntryRecord(t.Context(), kv, &handlers_eks.AccessEntryRecord{
 		ClusterName:        "alpha",
 		PrincipalARN:       testARN,
 		KubernetesUsername: testARN,

spinifex/gateway/oidc_discovery.go

@@ -1,6 +1,7 @@
 package gateway
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"log/slog"
@@ -8,7 +9,7 @@ import (
 
 	"github.com/go-chi/chi/v5"
 	handlers_eks "github.com/mulgadc/spinifex/spinifex/handlers/eks"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // oidcDiscoveryDocument is the subset of the OpenID Connect Discovery metadata
@@ -24,17 +25,17 @@ type oidcDiscoveryDocument struct {
 }
 
 // eksAccountKV opens the read-only per-account EKS bucket. Returns
-// nats.ErrBucketNotFound (unwrapped) when the account has no clusters, which
+// jetstream.ErrBucketNotFound (unwrapped) when the account has no clusters, which
 // the OIDC handlers translate to 404.
-func (gw *GatewayConfig) eksAccountKV(accountID string) (nats.KeyValue, error) {
+func (gw *GatewayConfig) eksAccountKV(ctx context.Context, accountID string) (jetstream.KeyValue, error) {
 	if gw.NATSConn == nil {
 		return nil, errors.New("gateway NATS connection not initialised")
 	}
-	js, err := gw.NATSConn.JetStream()
+	js, err := jetstream.New(gw.NATSConn)
 	if err != nil {
 		return nil, err
 	}
-	return js.KeyValue(handlers_eks.AccountBucketName(accountID))
+	return js.KeyValue(ctx, handlers_eks.AccountBucketName(accountID))
 }
 
 // OIDCDiscoveryDocument serves the cluster's OpenID discovery document. The
@@ -45,9 +46,9 @@ func (gw *GatewayConfig) OIDCDiscoveryDocument(w http.ResponseWriter, r *http.Re
 	accountID := chi.URLParam(r, "accountID")
 	clusterName := chi.URLParam(r, "clusterName")
 
-	kv, err := gw.eksAccountKV(accountID)
+	kv, err := gw.eksAccountKV(r.Context(), accountID)
 	if err != nil {
-		if errors.Is(err, nats.ErrBucketNotFound) {
+		if errors.Is(err, jetstream.ErrBucketNotFound) {
 			http.NotFound(w, r)
 			return
 		}
@@ -56,7 +57,7 @@ func (gw *GatewayConfig) OIDCDiscoveryDocument(w http.ResponseWriter, r *http.Re
 		return
 	}
 
-	meta, err := handlers_eks.GetClusterMeta(kv, clusterName)
+	meta, err := handlers_eks.GetClusterMeta(r.Context(), kv, clusterName)
 	if err != nil || meta == nil || meta.OIDCIssuer == "" {
 		http.NotFound(w, r)
 		return
@@ -78,9 +79,9 @@ func (gw *GatewayConfig) OIDCJWKS(w http.ResponseWriter, r *http.Request) {
 	accountID := chi.URLParam(r, "accountID")
 	clusterName := chi.URLParam(r, "clusterName")
 
-	kv, err := gw.eksAccountKV(accountID)
+	kv, err := gw.eksAccountKV(r.Context(), accountID)
 	if err != nil {
-		if errors.Is(err, nats.ErrBucketNotFound) {
+		if errors.Is(err, jetstream.ErrBucketNotFound) {
 			http.NotFound(w, r)
 			return
 		}
@@ -89,7 +90,7 @@ func (gw *GatewayConfig) OIDCJWKS(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	entry, err := kv.Get(handlers_eks.OIDCJWKSKey(clusterName))
+	entry, err := kv.Get(r.Context(), handlers_eks.OIDCJWKSKey(clusterName))
 	if err != nil {
 		http.NotFound(w, r)
 		return

spinifex/gateway/oidc_discovery_test.go

@@ -20,20 +20,21 @@ const (
 // into the per-account EKS bucket, returning a router wired to the same NATS.
 func seedDiscoveryCluster(t *testing.T) http.Handler {
 	t.Helper()
-	_, nc, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 
-	kv, err := handlers_eks.GetOrCreateAccountBucket(js, testDiscAccount, 1)
+	kv, err := handlers_eks.GetOrCreateAccountBucket(t.Context(), js, testDiscAccount, 1)
 	if err != nil {
 		t.Fatalf("account bucket: %v", err)
 	}
-	if err := handlers_eks.PutClusterMeta(kv, &handlers_eks.ClusterMeta{
+	if err := handlers_eks.PutClusterMeta(t.Context(), kv, &handlers_eks.ClusterMeta{
 		Name:       testDiscCluster,
 		Status:     handlers_eks.ClusterStatusActive,
 		OIDCIssuer: testDiscIssuer,
 	}); err != nil {
 		t.Fatalf("put cluster meta: %v", err)
 	}
-	if _, err := kv.Put(handlers_eks.OIDCJWKSKey(testDiscCluster), []byte(`{"keys":[{"kty":"EC","kid":"abc"}]}`)); err != nil {
+	if _, err := kv.Put(t.Context(), handlers_eks.OIDCJWKSKey(testDiscCluster), []byte(`{"keys":[{"kty":"EC","kid":"abc"}]}`)); err != nil {
 		t.Fatalf("put jwks: %v", err)
 	}
 

spinifex/handlers/acm/service_impl.go

@@ -39,8 +39,8 @@ var _ ACMService = (*ACMServiceImpl)(nil)
 
 // NewACMServiceImplWithNATS builds an ACM service backed by a JetStream KV
 // store. cfg may be nil (tests); region then falls back to the default.
-func NewACMServiceImplWithNATS(cfg *config.Config, nc *nats.Conn) (*ACMServiceImpl, error) {
-	store, err := NewStore(nc)
+func NewACMServiceImplWithNATS(ctx context.Context, cfg *config.Config, nc *nats.Conn) (*ACMServiceImpl, error) {
+	store, err := NewStore(ctx, nc)
 	if err != nil {
 		return nil, err
 	}
@@ -79,7 +79,7 @@ func (s *ACMServiceImpl) ImportCertificate(ctx context.Context, input *acm.Impor
 		certArn = s.mintCertificateArn(accountID)
 	} else {
 		// Re-import: the ARN must already exist and belong to the caller.
-		existing, gErr := s.store.GetCert(certArn)
+		existing, gErr := s.store.GetCert(ctx, certArn)
 		if gErr != nil {
 			return nil, errors.New(awserrors.ErrorInternalError)
 		}
@@ -105,7 +105,7 @@ func (s *ACMServiceImpl) ImportCertificate(ctx context.Context, input *acm.Impor
 		ImportedAt:       time.Now().UTC(),
 		Tags:             tagsToMap(input.Tags),
 	}
-	if err := s.store.PutCert(rec); err != nil {
+	if err := s.store.PutCert(ctx, rec); err != nil {
 		slog.ErrorContext(ctx, "ImportCertificate: store failed", "err", err)
 		return nil, errors.New(awserrors.ErrorInternalError)
 	}
@@ -119,7 +119,7 @@ func (s *ACMServiceImpl) DescribeCertificate(ctx context.Context, input *acm.Des
 	if input == nil || aws.StringValue(input.CertificateArn) == "" {
 		return nil, errors.New(awserrors.ErrorACMInvalidArn)
 	}
-	rec, err := s.lookupOwned(aws.StringValue(input.CertificateArn), accountID)
+	rec, err := s.lookupOwned(ctx, aws.StringValue(input.CertificateArn), accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -128,7 +128,7 @@ func (s *ACMServiceImpl) DescribeCertificate(ctx context.Context, input *acm.Des
 
 // ListCertificates returns summaries for every cert owned by accountID.
 func (s *ACMServiceImpl) ListCertificates(ctx context.Context, input *acm.ListCertificatesInput, accountID string) (*acm.ListCertificatesOutput, error) {
-	recs, err := s.store.ListCerts(accountID)
+	recs, err := s.store.ListCerts(ctx, accountID)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInternalError)
 	}
@@ -154,10 +154,10 @@ func (s *ACMServiceImpl) DeleteCertificate(ctx context.Context, input *acm.Delet
 	if input == nil || aws.StringValue(input.CertificateArn) == "" {
 		return nil, errors.New(awserrors.ErrorACMInvalidArn)
 	}
-	if _, err := s.lookupOwned(aws.StringValue(input.CertificateArn), accountID); err != nil {
+	if _, err := s.lookupOwned(ctx, aws.StringValue(input.CertificateArn), accountID); err != nil {
 		return nil, err
 	}
-	if _, err := s.store.DeleteCert(aws.StringValue(input.CertificateArn)); err != nil {
+	if _, err := s.store.DeleteCert(ctx, aws.StringValue(input.CertificateArn)); err != nil {
 		return nil, errors.New(awserrors.ErrorInternalError)
 	}
 	return &acm.DeleteCertificateOutput{}, nil
@@ -165,8 +165,8 @@ func (s *ACMServiceImpl) DeleteCertificate(ctx context.Context, input *acm.Delet
 
 // lookupOwned fetches a cert by ARN, returning ResourceNotFound when absent or
 // owned by a different account (no cross-account disclosure).
-func (s *ACMServiceImpl) lookupOwned(certArn, accountID string) (*CertRecord, error) {
-	rec, err := s.store.GetCert(certArn)
+func (s *ACMServiceImpl) lookupOwned(ctx context.Context, certArn, accountID string) (*CertRecord, error) {
+	rec, err := s.store.GetCert(ctx, certArn)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInternalError)
 	}
@@ -209,7 +209,7 @@ func (s *ACMServiceImpl) ListTagsForCertificate(ctx context.Context, input *acm.
 	if input == nil || aws.StringValue(input.CertificateArn) == "" {
 		return nil, errors.New(awserrors.ErrorACMInvalidArn)
 	}
-	rec, err := s.lookupOwned(aws.StringValue(input.CertificateArn), accountID)
+	rec, err := s.lookupOwned(ctx, aws.StringValue(input.CertificateArn), accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -221,7 +221,7 @@ func (s *ACMServiceImpl) AddTagsToCertificate(ctx context.Context, input *acm.Ad
 	if input == nil || aws.StringValue(input.CertificateArn) == "" {
 		return nil, errors.New(awserrors.ErrorACMInvalidArn)
 	}
-	rec, err := s.lookupOwned(aws.StringValue(input.CertificateArn), accountID)
+	rec, err := s.lookupOwned(ctx, aws.StringValue(input.CertificateArn), accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -229,7 +229,7 @@ func (s *ACMServiceImpl) AddTagsToCertificate(ctx context.Context, input *acm.Ad
 		rec.Tags = map[string]string{}
 	}
 	maps.Copy(rec.Tags, tagsToMap(input.Tags))
-	if err := s.store.PutCert(rec); err != nil {
+	if err := s.store.PutCert(ctx, rec); err != nil {
 		return nil, errors.New(awserrors.ErrorInternalError)
 	}
 	return &acm.AddTagsToCertificateOutput{}, nil
@@ -242,7 +242,7 @@ func (s *ACMServiceImpl) RemoveTagsFromCertificate(ctx context.Context, input *a
 	if input == nil || aws.StringValue(input.CertificateArn) == "" {
 		return nil, errors.New(awserrors.ErrorACMInvalidArn)
 	}
-	rec, err := s.lookupOwned(aws.StringValue(input.CertificateArn), accountID)
+	rec, err := s.lookupOwned(ctx, aws.StringValue(input.CertificateArn), accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -256,7 +256,7 @@ func (s *ACMServiceImpl) RemoveTagsFromCertificate(ctx context.Context, input *a
 		}
 		delete(rec.Tags, key)
 	}
-	if err := s.store.PutCert(rec); err != nil {
+	if err := s.store.PutCert(ctx, rec); err != nil {
 		return nil, errors.New(awserrors.ErrorInternalError)
 	}
 	return &acm.RemoveTagsFromCertificateOutput{}, nil

spinifex/handlers/acm/service_impl_test.go

@@ -25,7 +25,7 @@ const testAccountID = "000000000001"
 func setupACMService(t *testing.T) *ACMServiceImpl {
 	t.Helper()
 	_, nc, _ := testutil.StartTestJetStream(t)
-	svc, err := NewACMServiceImplWithNATS(nil, nc)
+	svc, err := NewACMServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 	return svc
 }

spinifex/handlers/acm/store.go

@@ -1,6 +1,7 @@
 package handlers_acm
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -8,8 +9,9 @@ import (
 	"strings"
 	"time"
 
-	"github.com/mulgadc/spinifex/spinifex/utils"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 const (
@@ -42,19 +44,20 @@ type CertRecord struct {
 
 // Store provides CRUD for ACM certificate records backed by JetStream KV.
 type Store struct {
-	kv nats.KeyValue
+	kv jetstream.KeyValue
 }
 
-// NewStore creates an ACM store using the provided NATS connection.
-func NewStore(nc *nats.Conn) (*Store, error) {
-	js, err := nc.JetStream()
+// NewStore creates an ACM store using the provided NATS connection. ctx bounds
+// the bucket get-or-create only; each operation carries its own.
+func NewStore(ctx context.Context, nc *nats.Conn) (*Store, error) {
+	js, err := jetstream.New(nc)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get JetStream context: %w", err)
 	}
 
-	kv, err := utils.GetOrCreateKVBucket(js, KVBucketACM, KVBucketACMVersion)
+	kv, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketACM, KVBucketACMVersion)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketACM, err)
+		return nil, err
 	}
 
 	slog.Info("ACM store initialized", "bucket", KVBucketACM)
@@ -71,20 +74,20 @@ func certKey(certArn string) string {
 }
 
 // PutCert stores (or replaces) a certificate record.
-func (s *Store) PutCert(rec *CertRecord) error {
+func (s *Store) PutCert(ctx context.Context, rec *CertRecord) error {
 	data, err := json.Marshal(rec)
 	if err != nil {
 		return fmt.Errorf("marshal cert: %w", err)
 	}
-	_, err = s.kv.Put(certKey(rec.CertificateArn), data)
+	_, err = s.kv.Put(ctx, certKey(rec.CertificateArn), data)
 	return err
 }
 
 // GetCert retrieves a certificate by ARN, returning (nil, nil) when absent.
-func (s *Store) GetCert(certArn string) (*CertRecord, error) {
-	entry, err := s.kv.Get(certKey(certArn))
+func (s *Store) GetCert(ctx context.Context, certArn string) (*CertRecord, error) {
+	entry, err := s.kv.Get(ctx, certKey(certArn))
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return nil, nil
 		}
 		return nil, err
@@ -97,24 +100,24 @@ func (s *Store) GetCert(certArn string) (*CertRecord, error) {
 }
 
 // DeleteCert removes a certificate by ARN. Returns (false, nil) when absent.
-func (s *Store) DeleteCert(certArn string) (bool, error) {
-	if _, err := s.kv.Get(certKey(certArn)); err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+func (s *Store) DeleteCert(ctx context.Context, certArn string) (bool, error) {
+	if _, err := s.kv.Get(ctx, certKey(certArn)); err != nil {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return false, nil
 		}
 		return false, err
 	}
-	if err := s.kv.Delete(certKey(certArn)); err != nil {
+	if err := s.kv.Delete(ctx, certKey(certArn)); err != nil {
 		return false, err
 	}
 	return true, nil
 }
 
 // ListCerts returns all certificate records owned by accountID.
-func (s *Store) ListCerts(accountID string) ([]*CertRecord, error) {
-	keys, err := s.kv.Keys()
+func (s *Store) ListCerts(ctx context.Context, accountID string) ([]*CertRecord, error) {
+	keys, err := s.kv.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, nil
 		}
 		return nil, err
@@ -124,7 +127,7 @@ func (s *Store) ListCerts(accountID string) ([]*CertRecord, error) {
 		if !strings.HasPrefix(key, KeyPrefixCert) {
 			continue
 		}
-		entry, err := s.kv.Get(key)
+		entry, err := s.kv.Get(ctx, key)
 		if err != nil {
 			continue
 		}

spinifex/handlers/ec2/account/service_impl.go

@@ -10,9 +10,11 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/mulgadc/spinifex/spinifex/config"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 const (
@@ -31,24 +33,24 @@ type AccountSettingsRecord struct {
 // AccountSettingsServiceImpl implements account settings operations with NATS JetStream persistence.
 type AccountSettingsServiceImpl struct {
 	config     *config.Config
-	js         nats.JetStreamContext
-	settingsKV nats.KeyValue
+	js         jetstream.JetStream
+	settingsKV jetstream.KeyValue
 }
 
 var _ AccountSettingsService = (*AccountSettingsServiceImpl)(nil)
 
 // NewAccountSettingsServiceImplWithNATS creates an account settings service with NATS JetStream for persistence.
-func NewAccountSettingsServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*AccountSettingsServiceImpl, error) {
-	js, err := natsConn.JetStream()
+func NewAccountSettingsServiceImplWithNATS(ctx context.Context, cfg *config.Config, natsConn *nats.Conn) (*AccountSettingsServiceImpl, error) {
+	js, err := jetstream.New(natsConn)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get JetStream context: %w", err)
 	}
 
-	settingsKV, err := utils.GetOrCreateKVBucket(js, KVBucketAccountSettings, 10)
+	settingsKV, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketAccountSettings, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create account settings KV bucket: %w", err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketAccountSettings, settingsKV, KVBucketAccountSettingsVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketAccountSettings, settingsKV, KVBucketAccountSettingsVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketAccountSettings, err)
 	}
 
@@ -71,10 +73,10 @@ func settingsKey(accountID string) string {
 }
 
 // getSettings retrieves current account settings.
-func (s *AccountSettingsServiceImpl) getSettings(accountID string) (*AccountSettingsRecord, error) {
-	entry, err := s.settingsKV.Get(settingsKey(accountID))
+func (s *AccountSettingsServiceImpl) getSettings(ctx context.Context, accountID string) (*AccountSettingsRecord, error) {
+	entry, err := s.settingsKV.Get(ctx, settingsKey(accountID))
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return &AccountSettingsRecord{
 				EbsEncryptionByDefault: false,
 				SerialConsoleAccess:    false,
@@ -92,13 +94,13 @@ func (s *AccountSettingsServiceImpl) getSettings(accountID string) (*AccountSett
 }
 
 // saveSettings saves current account settings.
-func (s *AccountSettingsServiceImpl) saveSettings(settings *AccountSettingsRecord, accountID string) error {
+func (s *AccountSettingsServiceImpl) saveSettings(ctx context.Context, settings *AccountSettingsRecord, accountID string) error {
 	data, err := json.Marshal(settings)
 	if err != nil {
 		return err
 	}
 
-	_, err = s.settingsKV.Put(settingsKey(accountID), data)
+	_, err = s.settingsKV.Put(ctx, settingsKey(accountID), data)
 	return err
 }
 
@@ -106,14 +108,14 @@ func (s *AccountSettingsServiceImpl) saveSettings(settings *AccountSettingsRecor
 func (s *AccountSettingsServiceImpl) EnableEbsEncryptionByDefault(ctx context.Context, input *ec2.EnableEbsEncryptionByDefaultInput, accountID string) (*ec2.EnableEbsEncryptionByDefaultOutput, error) {
 	slog.InfoContext(ctx, "EnableEbsEncryptionByDefault called", "accountID", accountID)
 
-	settings, err := s.getSettings(accountID)
+	settings, err := s.getSettings(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
 	settings.EbsEncryptionByDefault = true
 
-	if err := s.saveSettings(settings, accountID); err != nil {
+	if err := s.saveSettings(ctx, settings, accountID); err != nil {
 		return nil, err
 	}
 
@@ -126,14 +128,14 @@ func (s *AccountSettingsServiceImpl) EnableEbsEncryptionByDefault(ctx context.Co
 func (s *AccountSettingsServiceImpl) DisableEbsEncryptionByDefault(ctx context.Context, input *ec2.DisableEbsEncryptionByDefaultInput, accountID string) (*ec2.DisableEbsEncryptionByDefaultOutput, error) {
 	slog.InfoContext(ctx, "DisableEbsEncryptionByDefault called", "accountID", accountID)
 
-	settings, err := s.getSettings(accountID)
+	settings, err := s.getSettings(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
 	settings.EbsEncryptionByDefault = false
 
-	if err := s.saveSettings(settings, accountID); err != nil {
+	if err := s.saveSettings(ctx, settings, accountID); err != nil {
 		return nil, err
 	}
 
@@ -146,7 +148,7 @@ func (s *AccountSettingsServiceImpl) DisableEbsEncryptionByDefault(ctx context.C
 func (s *AccountSettingsServiceImpl) GetEbsEncryptionByDefault(ctx context.Context, input *ec2.GetEbsEncryptionByDefaultInput, accountID string) (*ec2.GetEbsEncryptionByDefaultOutput, error) {
 	slog.InfoContext(ctx, "GetEbsEncryptionByDefault called", "accountID", accountID)
 
-	settings, err := s.getSettings(accountID)
+	settings, err := s.getSettings(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -160,7 +162,7 @@ func (s *AccountSettingsServiceImpl) GetEbsEncryptionByDefault(ctx context.Conte
 func (s *AccountSettingsServiceImpl) GetSerialConsoleAccessStatus(ctx context.Context, input *ec2.GetSerialConsoleAccessStatusInput, accountID string) (*ec2.GetSerialConsoleAccessStatusOutput, error) {
 	slog.InfoContext(ctx, "GetSerialConsoleAccessStatus called", "accountID", accountID)
 
-	settings, err := s.getSettings(accountID)
+	settings, err := s.getSettings(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -174,14 +176,14 @@ func (s *AccountSettingsServiceImpl) GetSerialConsoleAccessStatus(ctx context.Co
 func (s *AccountSettingsServiceImpl) EnableSerialConsoleAccess(ctx context.Context, input *ec2.EnableSerialConsoleAccessInput, accountID string) (*ec2.EnableSerialConsoleAccessOutput, error) {
 	slog.InfoContext(ctx, "EnableSerialConsoleAccess called", "accountID", accountID)
 
-	settings, err := s.getSettings(accountID)
+	settings, err := s.getSettings(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
 	settings.SerialConsoleAccess = true
 
-	if err := s.saveSettings(settings, accountID); err != nil {
+	if err := s.saveSettings(ctx, settings, accountID); err != nil {
 		return nil, err
 	}
 
@@ -194,14 +196,14 @@ func (s *AccountSettingsServiceImpl) EnableSerialConsoleAccess(ctx context.Conte
 func (s *AccountSettingsServiceImpl) DisableSerialConsoleAccess(ctx context.Context, input *ec2.DisableSerialConsoleAccessInput, accountID string) (*ec2.DisableSerialConsoleAccessOutput, error) {
 	slog.InfoContext(ctx, "DisableSerialConsoleAccess called", "accountID", accountID)
 
-	settings, err := s.getSettings(accountID)
+	settings, err := s.getSettings(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
 	settings.SerialConsoleAccess = false
 
-	if err := s.saveSettings(settings, accountID); err != nil {
+	if err := s.saveSettings(ctx, settings, accountID); err != nil {
 		return nil, err
 	}
 

spinifex/handlers/ec2/account/service_impl_test.go

@@ -17,7 +17,7 @@ import (
 func TestNewAccountSettingsServiceImplWithNATS_NoJetStream(t *testing.T) {
 	_, nc := testutil.StartTestNATS(t)
 
-	svc, err := NewAccountSettingsServiceImplWithNATS(nil, nc)
+	svc, err := NewAccountSettingsServiceImplWithNATS(t.Context(), nil, nc)
 	assert.Error(t, err)
 	assert.Nil(t, svc)
 }
@@ -28,7 +28,7 @@ func setupTestAccountService(t *testing.T) *AccountSettingsServiceImpl {
 	t.Helper()
 	_, nc, _ := testutil.StartTestJetStream(t)
 
-	svc, err := NewAccountSettingsServiceImplWithNATS(nil, nc)
+	svc, err := NewAccountSettingsServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 	return svc
 }

spinifex/handlers/ec2/eigw/service_impl.go

@@ -15,9 +15,11 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
 	handlers_ec2_vpc "github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Ensure EgressOnlyIGWServiceImpl implements EgressOnlyIGWService.
@@ -40,27 +42,27 @@ type EgressOnlyIGWRecord struct {
 // EgressOnlyIGWServiceImpl implements Egress-only Internet Gateway operations with NATS JetStream persistence.
 type EgressOnlyIGWServiceImpl struct {
 	config *config.Config
-	eigwKV nats.KeyValue
-	vpcKV  nats.KeyValue
+	eigwKV jetstream.KeyValue
+	vpcKV  jetstream.KeyValue
 }
 
 // NewEgressOnlyIGWServiceImplWithNATS creates an Egress-only Internet Gateway service with NATS JetStream for persistence.
-func NewEgressOnlyIGWServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*EgressOnlyIGWServiceImpl, error) {
-	js, err := natsConn.JetStream()
+func NewEgressOnlyIGWServiceImplWithNATS(ctx context.Context, cfg *config.Config, natsConn *nats.Conn) (*EgressOnlyIGWServiceImpl, error) {
+	js, err := jetstream.New(natsConn)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get JetStream context: %w", err)
 	}
 
-	eigwKV, err := utils.GetOrCreateKVBucket(js, KVBucketEgressOnlyIGW, 10)
+	eigwKV, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketEgressOnlyIGW, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketEgressOnlyIGW, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketEgressOnlyIGW, eigwKV, KVBucketEgressOnlyIGWVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketEgressOnlyIGW, eigwKV, KVBucketEgressOnlyIGWVersion); err != nil {
 		return nil, fmt.Errorf("migrate %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)
+	vpcKV, err := kvutil.GetOrCreateBucket(ctx, js, handlers_ec2_vpc.KVBucketVPCs, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get VPC KV bucket: %w", err)
 	}
@@ -85,7 +87,7 @@ func (s *EgressOnlyIGWServiceImpl) CreateEgressOnlyInternetGateway(ctx context.C
 		slog.ErrorContext(ctx, "VPC KV unavailable, cannot verify VPC ownership")
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.vpcKV.Get(utils.AccountKey(accountID, *input.VpcId)); err != nil {
+	if _, err := s.vpcKV.Get(ctx, utils.AccountKey(accountID, *input.VpcId)); err != nil {
 		slog.WarnContext(ctx, "CreateEgressOnlyInternetGateway: VPC not found for account", "vpcId", *input.VpcId, "accountID", accountID)
 		return nil, errors.New(awserrors.ErrorInvalidVpcIDNotFound)
 	}
@@ -104,7 +106,7 @@ func (s *EgressOnlyIGWServiceImpl) CreateEgressOnlyInternetGateway(ctx context.C
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal Egress-only IGW record: %w", err)
 	}
-	if _, err := s.eigwKV.Put(utils.AccountKey(accountID, eigwID), data); err != nil {
+	if _, err := s.eigwKV.Put(ctx, utils.AccountKey(accountID, eigwID), data); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -125,13 +127,13 @@ func (s *EgressOnlyIGWServiceImpl) DeleteEgressOnlyInternetGateway(ctx context.C
 	key := utils.AccountKey(accountID, eigwID)
 
 	// Verify the EIGW exists before deleting
-	if _, err := s.eigwKV.Get(key); err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+	if _, err := s.eigwKV.Get(ctx, key); err != nil {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return nil, errors.New(awserrors.ErrorInvalidEgressOnlyInternetGatewayIdNotFound)
 		}
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
-	if err := s.eigwKV.Delete(key); err != nil {
+	if err := s.eigwKV.Delete(ctx, key); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -165,8 +167,8 @@ func (s *EgressOnlyIGWServiceImpl) DescribeEgressOnlyInternetGateways(ctx contex
 	}
 
 	prefix := accountID + "."
-	keys, err := s.eigwKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	keys, err := s.eigwKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -178,7 +180,7 @@ func (s *EgressOnlyIGWServiceImpl) DescribeEgressOnlyInternetGateways(ctx contex
 			continue
 		}
 
-		entry, err := s.eigwKV.Get(key)
+		entry, err := s.eigwKV.Get(ctx, key)
 		if err != nil {
 			slog.WarnContext(ctx, "Failed to get Egress-only IGW record", "key", key, "error", err)
 			continue
@@ -254,7 +256,7 @@ func (s *EgressOnlyIGWServiceImpl) ApplyRecordTags(input *ec2.CreateTagsInput, a
 	if input == nil {
 		return nil
 	}
-	return utils.MirrorKVRecordTags(s.eigwKV, accountID, "eigw-", input.Resources,
+	return utils.MirrorKVRecordTags(context.Background(), s.eigwKV, accountID, "eigw-", input.Resources,
 		func(r *EgressOnlyIGWRecord) *map[string]string { return &r.Tags },
 		utils.MergeTagsMut(input))
 }
@@ -265,7 +267,7 @@ func (s *EgressOnlyIGWServiceImpl) RemoveRecordTags(input *ec2.DeleteTagsInput,
 	if input == nil {
 		return nil
 	}
-	return utils.MirrorKVRecordTags(s.eigwKV, accountID, "eigw-", input.Resources,
+	return utils.MirrorKVRecordTags(context.Background(), s.eigwKV, accountID, "eigw-", input.Resources,
 		func(r *EgressOnlyIGWRecord) *map[string]string { return &r.Tags },
 		utils.RemoveTagsMut(input))
 }

spinifex/handlers/ec2/eigw/service_impl_test.go

@@ -26,7 +26,7 @@ func setupTestEIGWService(t *testing.T) *EgressOnlyIGWServiceImpl {
 	}
 	testutil.SeedKV(t, js, handlers_ec2_vpc.KVBucketVPCs, vpcEntries)
 
-	svc, err := NewEgressOnlyIGWServiceImplWithNATS(nil, nc)
+	svc, err := NewEgressOnlyIGWServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 	return svc
 }
@@ -285,7 +285,7 @@ func TestCreateEgressOnlyInternetGateway_CrossAccountVPCRejected(t *testing.T) {
 	require.NoError(t, err)
 
 	// Create EIGW service (VPC KV will be populated since we created it before EIGW init)
-	svc, err := NewEgressOnlyIGWServiceImplWithNATS(nil, nc)
+	svc, err := NewEgressOnlyIGWServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 
 	// Other account tries to create EIGW in testAccountID's VPC — should fail

spinifex/handlers/ec2/eip/service_impl.go

@@ -14,10 +14,12 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
 	handlers_ec2_vpc "github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/network/topology"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Ensure EIPServiceImpl implements EIPService.
@@ -25,7 +27,7 @@ var _ EIPService = (*EIPServiceImpl)(nil)
 
 // EIPServiceImpl implements Elastic IP operations with NATS JetStream persistence.
 type EIPServiceImpl struct {
-	eipKV        nats.KeyValue
+	eipKV        jetstream.KeyValue
 	externalIPAM *handlers_ec2_vpc.ExternalIPAM
 	vpcService   handlers_ec2_vpc.VPCService
 	natsConn     *nats.Conn
@@ -41,17 +43,17 @@ type natEvent struct {
 }
 
 // NewEIPServiceImpl creates a new EIP service backed by NATS JetStream KV.
-func NewEIPServiceImpl(natsConn *nats.Conn, externalIPAM *handlers_ec2_vpc.ExternalIPAM, vpcService handlers_ec2_vpc.VPCService) (*EIPServiceImpl, error) {
-	js, err := natsConn.JetStream()
+func NewEIPServiceImpl(ctx context.Context, natsConn *nats.Conn, externalIPAM *handlers_ec2_vpc.ExternalIPAM, vpcService handlers_ec2_vpc.VPCService) (*EIPServiceImpl, error) {
+	js, err := jetstream.New(natsConn)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get JetStream context: %w", err)
 	}
 
-	eipKV, err := utils.GetOrCreateKVBucket(js, KVBucketEIPs, 10)
+	eipKV, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketEIPs, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketEIPs, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketEIPs, eipKV, KVBucketEIPsVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketEIPs, eipKV, KVBucketEIPsVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketEIPs, err)
 	}
 
@@ -104,7 +106,7 @@ func (s *EIPServiceImpl) AllocateAddress(ctx context.Context, input *ec2.Allocat
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal EIP record: %w", err)
 	}
-	if _, err := s.eipKV.Put(utils.AccountKey(accountID, allocID), data); err != nil {
+	if _, err := s.eipKV.Put(ctx, utils.AccountKey(accountID, allocID), data); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -127,7 +129,7 @@ func (s *EIPServiceImpl) ReleaseAddress(ctx context.Context, input *ec2.ReleaseA
 	allocID := *input.AllocationId
 	key := utils.AccountKey(accountID, allocID)
 
-	entry, err := s.eipKV.Get(key)
+	entry, err := s.eipKV.Get(ctx, key)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidAllocationIDNotFound)
 	}
@@ -148,7 +150,7 @@ func (s *EIPServiceImpl) ReleaseAddress(ctx context.Context, input *ec2.ReleaseA
 		slog.WarnContext(ctx, "Failed to release IP back to IPAM pool", "allocationId", allocID, "ip", record.PublicIp, "pool", record.PoolName, "err", err)
 	}
 
-	if err := s.eipKV.Delete(key); err != nil {
+	if err := s.eipKV.Delete(ctx, key); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -166,7 +168,7 @@ func (s *EIPServiceImpl) AssociateAddress(ctx context.Context, input *ec2.Associ
 	allocID := *input.AllocationId
 	key := utils.AccountKey(accountID, allocID)
 
-	entry, err := s.eipKV.Get(key)
+	entry, err := s.eipKV.Get(ctx, key)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidAllocationIDNotFound)
 	}
@@ -222,7 +224,7 @@ func (s *EIPServiceImpl) AssociateAddress(ctx context.Context, input *ec2.Associ
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal EIP record: %w", err)
 	}
-	if _, err := s.eipKV.Update(key, data, entry.Revision()); err != nil {
+	if _, err := s.eipKV.Update(ctx, key, data, entry.Revision()); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -252,7 +254,7 @@ func (s *EIPServiceImpl) DisassociateAddress(ctx context.Context, input *ec2.Dis
 	associationID := *input.AssociationId
 
 	// Find the EIP record by association ID.
-	record, key, revision, err := s.findByAssociationID(accountID, associationID)
+	record, key, revision, err := s.findByAssociationID(ctx, accountID, associationID)
 	if err != nil {
 		return nil, err
 	}
@@ -280,7 +282,7 @@ func (s *EIPServiceImpl) DisassociateAddress(ctx context.Context, input *ec2.Dis
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal EIP record: %w", err)
 	}
-	if _, err := s.eipKV.Update(key, data, revision); err != nil {
+	if _, err := s.eipKV.Update(ctx, key, data, revision); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -321,8 +323,8 @@ func (s *EIPServiceImpl) DescribeAddresses(ctx context.Context, input *ec2.Descr
 	}
 
 	prefix := accountID + "."
-	keys, err := s.eipKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	keys, err := s.eipKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -335,7 +337,7 @@ func (s *EIPServiceImpl) DescribeAddresses(ctx context.Context, input *ec2.Descr
 			continue
 		}
 
-		entry, err := s.eipKV.Get(k)
+		entry, err := s.eipKV.Get(ctx, k)
 		if err != nil {
 			slog.WarnContext(ctx, "Failed to get EIP record", "key", k, "error", err)
 			continue
@@ -426,7 +428,7 @@ func (s *EIPServiceImpl) DescribeAddressesAttribute(ctx context.Context, input *
 				continue
 			}
 			key := utils.AccountKey(accountID, *id)
-			entry, err := s.eipKV.Get(key)
+			entry, err := s.eipKV.Get(ctx, key)
 			if err != nil {
 				continue // not found, skip
 			}
@@ -443,8 +445,8 @@ func (s *EIPServiceImpl) DescribeAddressesAttribute(ctx context.Context, input *
 	} else {
 		// No filter — scan all EIPs for this account.
 		prefix := accountID + "."
-		keys, err := s.eipKV.Keys()
-		if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+		keys, err := s.eipKV.Keys(ctx)
+		if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, errors.New(awserrors.ErrorServerInternal)
 		}
 		for _, k := range keys {
@@ -454,7 +456,7 @@ func (s *EIPServiceImpl) DescribeAddressesAttribute(ctx context.Context, input *
 			if !strings.HasPrefix(k, prefix) {
 				continue
 			}
-			entry, err := s.eipKV.Get(k)
+			entry, err := s.eipKV.Get(ctx, k)
 			if err != nil {
 				slog.WarnContext(ctx, "Failed to get EIP record", "key", k, "error", err)
 				continue
@@ -537,9 +539,9 @@ func (s *EIPServiceImpl) lookupENIByInstance(ctx context.Context, accountID, ins
 }
 
 // findByAssociationID scans EIP records to find one matching the given association ID.
-func (s *EIPServiceImpl) findByAssociationID(accountID, associationID string) (*EIPRecord, string, uint64, error) {
+func (s *EIPServiceImpl) findByAssociationID(ctx context.Context, accountID, associationID string) (*EIPRecord, string, uint64, error) {
 	prefix := accountID + "."
-	keys, err := s.eipKV.Keys()
+	keys, err := s.eipKV.Keys(ctx)
 	if err != nil {
 		return nil, "", 0, errors.New(awserrors.ErrorServerInternal)
 	}
@@ -552,7 +554,7 @@ func (s *EIPServiceImpl) findByAssociationID(accountID, associationID string) (*
 			continue
 		}
 
-		entry, err := s.eipKV.Get(k)
+		entry, err := s.eipKV.Get(ctx, k)
 		if err != nil {
 			continue
 		}
@@ -578,7 +580,7 @@ func (s *EIPServiceImpl) AssociatedPublicIPForInstance(ctx context.Context, acco
 		return "", false
 	}
 	prefix := accountID + "."
-	keys, err := s.eipKV.Keys()
+	keys, err := s.eipKV.Keys(ctx)
 	if err != nil {
 		return "", false
 	}
@@ -586,7 +588,7 @@ func (s *EIPServiceImpl) AssociatedPublicIPForInstance(ctx context.Context, acco
 		if k == utils.VersionKey || !strings.HasPrefix(k, prefix) {
 			continue
 		}
-		entry, err := s.eipKV.Get(k)
+		entry, err := s.eipKV.Get(ctx, k)
 		if err != nil {
 			continue
 		}
@@ -607,12 +609,13 @@ func (s *EIPServiceImpl) AssociatedPublicIPForInstance(ctx context.Context, acco
 // gone from the manager the per-instance release never runs, so an
 // internet-facing ALB's EIP would otherwise orphan. No-op when nothing matches.
 func (s *EIPServiceImpl) ReleaseAddressByInstanceID(instanceID string) error {
+	ctx := context.Background()
 	if instanceID == "" {
 		return nil
 	}
-	keys, err := s.eipKV.Keys()
+	keys, err := s.eipKV.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil
 		}
 		return fmt.Errorf("list eip keys: %w", err)
@@ -622,7 +625,7 @@ func (s *EIPServiceImpl) ReleaseAddressByInstanceID(instanceID string) error {
 		if k == utils.VersionKey {
 			continue
 		}
-		entry, err := s.eipKV.Get(k)
+		entry, err := s.eipKV.Get(ctx, k)
 		if err != nil {
 			continue
 		}
@@ -639,13 +642,13 @@ func (s *EIPServiceImpl) ReleaseAddressByInstanceID(instanceID string) error {
 			continue
 		}
 		if record.AssociationId != "" {
-			if _, err := s.DisassociateAddress(context.Background(), &ec2.DisassociateAddressInput{
+			if _, err := s.DisassociateAddress(ctx, &ec2.DisassociateAddressInput{
 				AssociationId: aws.String(record.AssociationId),
 			}, accountID); err != nil {
 				errs = append(errs, fmt.Errorf("disassociate %s: %w", record.AllocationId, err))
 			}
 		}
-		if _, err := s.ReleaseAddress(context.Background(), &ec2.ReleaseAddressInput{
+		if _, err := s.ReleaseAddress(ctx, &ec2.ReleaseAddressInput{
 			AllocationId: aws.String(record.AllocationId),
 		}, accountID); err != nil {
 			errs = append(errs, fmt.Errorf("release %s: %w", record.AllocationId, err))
@@ -701,7 +704,7 @@ func (s *EIPServiceImpl) ApplyRecordTags(input *ec2.CreateTagsInput, accountID s
 	if input == nil {
 		return nil
 	}
-	return utils.MirrorKVRecordTags(s.eipKV, accountID, "eipalloc-", input.Resources,
+	return utils.MirrorKVRecordTags(context.Background(), s.eipKV, accountID, "eipalloc-", input.Resources,
 		func(r *EIPRecord) *map[string]string { return &r.Tags },
 		utils.MergeTagsMut(input))
 }
@@ -712,7 +715,7 @@ func (s *EIPServiceImpl) RemoveRecordTags(input *ec2.DeleteTagsInput, accountID
 	if input == nil {
 		return nil
 	}
-	return utils.MirrorKVRecordTags(s.eipKV, accountID, "eipalloc-", input.Resources,
+	return utils.MirrorKVRecordTags(context.Background(), s.eipKV, accountID, "eipalloc-", input.Resources,
 		func(r *EIPRecord) *map[string]string { return &r.Tags },
 		utils.RemoveTagsMut(input))
 }

spinifex/handlers/ec2/eip/service_impl_test.go

@@ -29,13 +29,14 @@ func testPool() external.ExternalPoolConfig {
 
 func setupTestEIP(t *testing.T) (*EIPServiceImpl, *handlers_ec2_vpc.ExternalIPAM) {
 	t.Helper()
-	_, nc, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 
 	pool := testPool()
-	ipam, err := handlers_ec2_vpc.NewExternalIPAM(js, []external.ExternalPoolConfig{pool})
+	ipam, err := handlers_ec2_vpc.NewExternalIPAM(t.Context(), js, []external.ExternalPoolConfig{pool})
 	require.NoError(t, err)
 
-	svc, err := NewEIPServiceImpl(nc, ipam, nil)
+	svc, err := NewEIPServiceImpl(t.Context(), nc, ipam, nil)
 	require.NoError(t, err)
 
 	return svc, ipam
@@ -118,7 +119,7 @@ func TestEIP_ReleaseWhileAssociated(t *testing.T) {
 	// To test the associated path, we need to manipulate the KV directly.
 
 	// Get the KV entry and update state to "associated"
-	entry, err := svc.eipKV.Get(testAccountID + "." + allocID)
+	entry, err := svc.eipKV.Get(t.Context(), testAccountID+"."+allocID)
 	require.NoError(t, err)
 
 	var record EIPRecord
@@ -130,7 +131,7 @@ func TestEIP_ReleaseWhileAssociated(t *testing.T) {
 
 	data, err := json.Marshal(record)
 	require.NoError(t, err)
-	_, err = svc.eipKV.Update(testAccountID+"."+allocID, data, entry.Revision())
+	_, err = svc.eipKV.Update(t.Context(), testAccountID+"."+allocID, data, entry.Revision())
 	require.NoError(t, err)
 
 	// Now try to release — should fail
@@ -150,7 +151,7 @@ func TestEIP_ReleaseByInstanceID_ReclaimsAssociatedEIP(t *testing.T) {
 
 	// Mark the EIP associated with a system instance, as an internet-facing ALB
 	// leaves it. ENIId stays empty so disassociate skips the VPC-service lookup.
-	entry, err := svc.eipKV.Get(testAccountID + "." + allocID)
+	entry, err := svc.eipKV.Get(t.Context(), testAccountID+"."+allocID)
 	require.NoError(t, err)
 	var record EIPRecord
 	require.NoError(t, json.Unmarshal(entry.Value(), &record))
@@ -159,7 +160,7 @@ func TestEIP_ReleaseByInstanceID_ReclaimsAssociatedEIP(t *testing.T) {
 	record.InstanceId = "i-alb-test"
 	data, err := json.Marshal(record)
 	require.NoError(t, err)
-	_, err = svc.eipKV.Update(testAccountID+"."+allocID, data, entry.Revision())
+	_, err = svc.eipKV.Update(t.Context(), testAccountID+"."+allocID, data, entry.Revision())
 	require.NoError(t, err)
 
 	// Backstop release by instance disassociates then frees the allocation.
@@ -180,7 +181,7 @@ func TestEIP_ReleaseByInstanceID_NoMatchIsNoOp(t *testing.T) {
 	// An instance with no recorded EIP must leave existing allocations untouched.
 	require.NoError(t, svc.ReleaseAddressByInstanceID("i-unrelated"))
 
-	_, err = svc.eipKV.Get(testAccountID + "." + allocID)
+	_, err = svc.eipKV.Get(t.Context(), testAccountID+"."+allocID)
 	assert.NoError(t, err)
 
 	// Empty instance ID is a no-op too.
@@ -316,7 +317,7 @@ func TestEIP_RecordToEC2_WithoutTags(t *testing.T) {
 func TestEIP_FindByAssociationID_NotFound(t *testing.T) {
 	svc, _ := setupTestEIP(t)
 
-	_, _, _, err := svc.findByAssociationID(testAccountID, "eipassoc-nonexistent")
+	_, _, _, err := svc.findByAssociationID(t.Context(), testAccountID, "eipassoc-nonexistent")
 	assert.Error(t, err)
 	assert.Contains(t, err.Error(), "InvalidAssociationID.NotFound")
 }
@@ -463,7 +464,7 @@ func TestEIP_DescribeAddresses_FilterByInstanceId(t *testing.T) {
 	require.NoError(t, err)
 
 	// Manually set instance-id on the record
-	entry, err := svc.eipKV.Get(testAccountID + "." + *out.AllocationId)
+	entry, err := svc.eipKV.Get(t.Context(), testAccountID+"."+*out.AllocationId)
 	require.NoError(t, err)
 	var record EIPRecord
 	require.NoError(t, json.Unmarshal(entry.Value(), &record))
@@ -471,7 +472,7 @@ func TestEIP_DescribeAddresses_FilterByInstanceId(t *testing.T) {
 	record.State = "associated"
 	data, err := json.Marshal(record)
 	require.NoError(t, err)
-	_, err = svc.eipKV.Update(testAccountID+"."+*out.AllocationId, data, entry.Revision())
+	_, err = svc.eipKV.Update(t.Context(), testAccountID+"."+*out.AllocationId, data, entry.Revision())
 	require.NoError(t, err)
 
 	// Allocate another without association

spinifex/handlers/ec2/igw/service_impl.go

@@ -15,10 +15,12 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
 	handlers_ec2_vpc "github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/types"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Ensure IGWServiceImpl implements IGWService.
@@ -48,8 +50,8 @@ type GatePublisher interface {
 // IGWServiceImpl implements Internet Gateway operations with NATS JetStream persistence.
 type IGWServiceImpl struct {
 	config        *config.Config
-	igwKV         nats.KeyValue
-	vpcKV         nats.KeyValue
+	igwKV         jetstream.KeyValue
+	vpcKV         jetstream.KeyValue
 	natsConn      *nats.Conn
 	gatePublisher GatePublisher
 }
@@ -61,22 +63,22 @@ func (s *IGWServiceImpl) SetGatePublisher(p GatePublisher) {
 }
 
 // NewIGWServiceImplWithNATS creates an Internet Gateway service with NATS JetStream for persistence.
-func NewIGWServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*IGWServiceImpl, error) {
-	js, err := natsConn.JetStream()
+func NewIGWServiceImplWithNATS(ctx context.Context, cfg *config.Config, natsConn *nats.Conn) (*IGWServiceImpl, error) {
+	js, err := jetstream.New(natsConn)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get JetStream context: %w", err)
 	}
 
-	igwKV, err := utils.GetOrCreateKVBucket(js, KVBucketIGW, 10)
+	igwKV, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketIGW, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketIGW, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketIGW, igwKV, KVBucketIGWVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketIGW, igwKV, KVBucketIGWVersion); err != nil {
 		return nil, fmt.Errorf("migrate %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)
+	vpcKV, err := kvutil.GetOrCreateBucket(ctx, js, handlers_ec2_vpc.KVBucketVPCs, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get VPC KV bucket: %w", err)
 	}
@@ -115,7 +117,7 @@ func (s *IGWServiceImpl) createIGW(ctx context.Context, input *ec2.CreateInterne
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal IGW record: %w", err)
 	}
-	if _, err := s.igwKV.Put(utils.AccountKey(accountID, igwID), data); err != nil {
+	if _, err := s.igwKV.Put(ctx, utils.AccountKey(accountID, igwID), data); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -135,12 +137,12 @@ func (s *IGWServiceImpl) DeleteInternetGateway(ctx context.Context, input *ec2.D
 	igwID := *input.InternetGatewayId
 	key := utils.AccountKey(accountID, igwID)
 
-	entry, err := s.igwKV.Get(key)
+	entry, err := s.igwKV.Get(ctx, key)
 	if err != nil {
 		// AWS-faithful: an absent internet gateway is NotFound (provider
 		// tolerates it on destroy); destroy orchestration tolerates it too.
 		// A transient read error stays a server error.
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return nil, errors.New(awserrors.ErrorInvalidInternetGatewayIDNotFound)
 		}
 		return nil, errors.New(awserrors.ErrorServerInternal)
@@ -156,7 +158,7 @@ func (s *IGWServiceImpl) DeleteInternetGateway(ctx context.Context, input *ec2.D
 		return nil, errors.New(awserrors.ErrorDependencyViolation)
 	}
 
-	if err := s.igwKV.Delete(key); err != nil {
+	if err := s.igwKV.Delete(ctx, key); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -190,8 +192,8 @@ func (s *IGWServiceImpl) DescribeInternetGateways(ctx context.Context, input *ec
 	}
 
 	prefix := accountID + "."
-	keys, err := s.igwKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	keys, err := s.igwKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -205,7 +207,7 @@ func (s *IGWServiceImpl) DescribeInternetGateways(ctx context.Context, input *ec
 			continue
 		}
 
-		entry, err := s.igwKV.Get(key)
+		entry, err := s.igwKV.Get(ctx, key)
 		if err != nil {
 			slog.WarnContext(ctx, "Failed to get IGW record", "key", key, "error", err)
 			continue
@@ -287,7 +289,7 @@ func (s *IGWServiceImpl) AttachInternetGateway(ctx context.Context, input *ec2.A
 	vpcID := *input.VpcId
 	key := utils.AccountKey(accountID, igwID)
 
-	entry, err := s.igwKV.Get(key)
+	entry, err := s.igwKV.Get(ctx, key)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidInternetGatewayIDNotFound)
 	}
@@ -306,7 +308,7 @@ func (s *IGWServiceImpl) AttachInternetGateway(ctx context.Context, input *ec2.A
 		slog.ErrorContext(ctx, "VPC KV unavailable, cannot verify VPC ownership")
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.vpcKV.Get(utils.AccountKey(accountID, vpcID)); err != nil {
+	if _, err := s.vpcKV.Get(ctx, utils.AccountKey(accountID, vpcID)); err != nil {
 		slog.WarnContext(ctx, "AttachInternetGateway: VPC not found for account", "vpcId", vpcID, "accountID", accountID)
 		return nil, errors.New(awserrors.ErrorInvalidVpcIDNotFound)
 	}
@@ -318,7 +320,7 @@ func (s *IGWServiceImpl) AttachInternetGateway(ctx context.Context, input *ec2.A
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.igwKV.Put(key, data); err != nil {
+	if _, err := s.igwKV.Put(ctx, key, data); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -358,7 +360,7 @@ func (s *IGWServiceImpl) DetachInternetGateway(ctx context.Context, input *ec2.D
 	vpcID := *input.VpcId
 	key := utils.AccountKey(accountID, igwID)
 
-	entry, err := s.igwKV.Get(key)
+	entry, err := s.igwKV.Get(ctx, key)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidInternetGatewayIDNotFound)
 	}
@@ -379,7 +381,7 @@ func (s *IGWServiceImpl) DetachInternetGateway(ctx context.Context, input *ec2.D
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.igwKV.Put(key, data); err != nil {
+	if _, err := s.igwKV.Put(ctx, key, data); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -436,7 +438,7 @@ func (s *IGWServiceImpl) ApplyRecordTags(input *ec2.CreateTagsInput, accountID s
 	if input == nil {
 		return nil
 	}
-	return utils.MirrorKVRecordTags(s.igwKV, accountID, "igw-", input.Resources,
+	return utils.MirrorKVRecordTags(context.Background(), s.igwKV, accountID, "igw-", input.Resources,
 		func(r *IGWRecord) *map[string]string { return &r.Tags },
 		utils.MergeTagsMut(input))
 }
@@ -447,7 +449,7 @@ func (s *IGWServiceImpl) RemoveRecordTags(input *ec2.DeleteTagsInput, accountID
 	if input == nil {
 		return nil
 	}
-	return utils.MirrorKVRecordTags(s.igwKV, accountID, "igw-", input.Resources,
+	return utils.MirrorKVRecordTags(context.Background(), s.igwKV, accountID, "igw-", input.Resources,
 		func(r *IGWRecord) *map[string]string { return &r.Tags },
 		utils.RemoveTagsMut(input))
 }

spinifex/handlers/ec2/igw/service_impl_test.go

@@ -56,7 +56,7 @@ func setupTestIGWService(t *testing.T) (*IGWServiceImpl, *nats.Conn) {
 	}
 	testutil.SeedKV(t, js, handlers_ec2_vpc.KVBucketVPCs, vpcEntries)
 
-	svc, err := NewIGWServiceImplWithNATS(nil, nc)
+	svc, err := NewIGWServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 	return svc, nc
 }
@@ -463,14 +463,13 @@ func TestCreateInternetGateway_PublishesNoEvent(t *testing.T) {
 func TestAttachInternetGateway_CrossAccountVPCRejected(t *testing.T) {
 	svc, nc := setupTestIGWService(t)
 
-	// Create VPC KV bucket and add a VPC owned by testAccountID
-	js, err := nc.JetStream()
-	require.NoError(t, err)
-	vpcKV, err := js.CreateKeyValue(&nats.KeyValueConfig{Bucket: handlers_ec2_vpc.KVBucketVPCs, History: 1})
+	// Add a VPC owned by testAccountID to the existing VPC bucket.
+	js := testutil.NewJetStream(t, nc)
+	vpcKV, err := js.KeyValue(t.Context(), handlers_ec2_vpc.KVBucketVPCs)
 	require.NoError(t, err)
 
 	vpcID := "vpc-alpha123"
-	_, err = vpcKV.Put(utils.AccountKey(testAccountID, vpcID), []byte(`{"vpc_id":"vpc-alpha123","state":"available"}`))
+	_, err = vpcKV.Put(t.Context(), utils.AccountKey(testAccountID, vpcID), []byte(`{"vpc_id":"vpc-alpha123","state":"available"}`))
 	require.NoError(t, err)
 
 	// Refresh service to pick up the VPC KV bucket

spinifex/handlers/ec2/instance/service.go

@@ -88,7 +88,7 @@ type StoppedInstanceStore interface {
 	// UpdateStoppedInstance atomically applies mutate to the current stopped
 	// record under optimistic concurrency (CAS) with createIfAbsent=false, so
 	// a caller racing a winning ClaimStoppedInstance gets a clean
-	// nats.ErrKeyNotFound instead of resurrecting a deleted record.
+	// jetstream.ErrKeyNotFound instead of resurrecting a deleted record.
 	UpdateStoppedInstance(instanceID string, mutate func(*vm.VM)) (*vm.VM, error)
 	WriteTerminatedInstance(instanceID string, instance *vm.VM) error
 	// ClaimStoppedInstance atomically removes instanceID's record and

spinifex/handlers/ec2/instance/service_impl.go

@@ -32,6 +32,7 @@ import (
 	"github.com/mulgadc/viperblock/viperblock"
 	"github.com/mulgadc/viperblock/viperblock/backends/s3"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // VolumeInfo holds volume information returned from GenerateVolumes
@@ -374,14 +375,14 @@ func (s *InstanceServiceImpl) TagStoppedInstance(ctx context.Context, instanceID
 	// CAS the tag mutation into the shared KV record instead of a wholesale
 	// WriteStoppedInstance: createIfAbsent=false means a claim that deletes
 	// the record between the Load above and this write fails cleanly with
-	// nats.ErrKeyNotFound rather than resurrecting a stale stopped entry.
+	// jetstream.ErrKeyNotFound rather than resurrecting a stale stopped entry.
 	newTags := instance.Instance.Tags
 	if _, err := s.stoppedStore.UpdateStoppedInstance(instanceID, func(v *vm.VM) {
 		if v.Instance != nil {
 			v.Instance.Tags = newTags
 		}
 	}); err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			slog.WarnContext(ctx, "TagStoppedInstance: instance claimed concurrently, not resurrecting", "instanceId", instanceID)
 			return errors.New(awserrors.ErrorInvalidInstanceIDNotFound)
 		}
@@ -2056,7 +2057,7 @@ func (s *InstanceServiceImpl) ModifyInstanceAttribute(ctx context.Context, input
 	// CAS the mutations into the shared KV record instead of a wholesale
 	// WriteStoppedInstance: createIfAbsent=false means a claim that deletes
 	// the record between the Load above and this write fails cleanly with
-	// nats.ErrKeyNotFound rather than resurrecting a stale stopped entry.
+	// jetstream.ErrKeyNotFound rather than resurrecting a stale stopped entry.
 	_, err = s.stoppedStore.UpdateStoppedInstance(instanceID, func(v *vm.VM) {
 		if input.InstanceType != nil && input.InstanceType.Value != nil && v.Instance != nil {
 			newType := *input.InstanceType.Value
@@ -2085,7 +2086,7 @@ func (s *InstanceServiceImpl) ModifyInstanceAttribute(ctx context.Context, input
 		}
 	})
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			// The record existed moments ago (Load above) but a concurrent
 			// claim removed it: the instance is mid-transition, not gone.
 			slog.WarnContext(ctx, "ModifyInstanceAttribute: instance claimed concurrently, not resurrecting", "instanceId", instanceID)

spinifex/handlers/ec2/instance/service_impl_test.go

@@ -22,6 +22,7 @@ import (
 	"github.com/mulgadc/viperblock/viperblock"
 	"github.com/nats-io/nats-server/v2/server"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -714,7 +715,7 @@ func (f *fakeStoppedStore) DeleteTerminatedInstance(string) error { return nil }
 
 // UpdateStoppedInstance mimics the real CAS semantics: mutate runs under the
 // store lock against loadByID's entry, and a missing record (e.g. removed by
-// a concurrent ClaimStoppedInstance) returns nats.ErrKeyNotFound instead of
+// a concurrent ClaimStoppedInstance) returns jetstream.ErrKeyNotFound instead of
 // resurrecting it — the same createIfAbsent=false contract as JetStreamManager.
 func (f *fakeStoppedStore) UpdateStoppedInstance(id string, mutate func(*vm.VM)) (*vm.VM, error) {
 	f.mu.Lock()
@@ -725,7 +726,7 @@ func (f *fakeStoppedStore) UpdateStoppedInstance(id string, mutate func(*vm.VM))
 	}
 	v, ok := f.loadByID[id]
 	if !ok {
-		return nil, nats.ErrKeyNotFound
+		return nil, jetstream.ErrKeyNotFound
 	}
 	mutate(v)
 	return v, nil

spinifex/handlers/ec2/launchtemplate/service_impl.go

@@ -17,9 +17,11 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Ensure LaunchTemplateServiceImpl implements LaunchTemplateService.
@@ -48,21 +50,21 @@ const (
 type LaunchTemplateServiceImpl struct {
 	config   *config.Config
 	natsConn *nats.Conn
-	kv       nats.KeyValue
+	kv       jetstream.KeyValue
 }
 
 // NewLaunchTemplateServiceImplWithNATS creates a launch template service with NATS JetStream.
-func NewLaunchTemplateServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*LaunchTemplateServiceImpl, error) {
-	js, err := natsConn.JetStream()
+func NewLaunchTemplateServiceImplWithNATS(ctx context.Context, cfg *config.Config, natsConn *nats.Conn) (*LaunchTemplateServiceImpl, error) {
+	js, err := jetstream.New(natsConn)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get JetStream context: %w", err)
 	}
 
-	kv, err := utils.GetOrCreateKVBucket(js, KVBucketLaunchTemplates, 10)
+	kv, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketLaunchTemplates, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketLaunchTemplates, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketLaunchTemplates, kv, KVBucketLaunchTemplatesVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketLaunchTemplates, kv, KVBucketLaunchTemplatesVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketLaunchTemplates, err)
 	}
 
@@ -102,8 +104,8 @@ func versionPrefix(accountID, ltID string) string {
 // --- record load helpers ---
 
 // getHeaderByID reads a header and its KV entry (for CAS) by launch template id.
-func (s *LaunchTemplateServiceImpl) getHeaderByID(accountID, ltID string) (*LaunchTemplateHeader, nats.KeyValueEntry, error) {
-	entry, err := s.kv.Get(headerKey(accountID, ltID))
+func (s *LaunchTemplateServiceImpl) getHeaderByID(ctx context.Context, accountID, ltID string) (*LaunchTemplateHeader, jetstream.KeyValueEntry, error) {
+	entry, err := s.kv.Get(ctx, headerKey(accountID, ltID))
 	if err != nil {
 		return nil, nil, errors.New(awserrors.ErrorInvalidLaunchTemplateIdNotFound)
 	}
@@ -116,13 +118,13 @@ func (s *LaunchTemplateServiceImpl) getHeaderByID(accountID, ltID string) (*Laun
 
 // getHeaderByName resolves a template name to its header via the name index.
 // A dangling name index (header missing) is treated as not-found.
-func (s *LaunchTemplateServiceImpl) getHeaderByName(accountID, name string) (*LaunchTemplateHeader, nats.KeyValueEntry, error) {
-	entry, err := s.kv.Get(nameKey(accountID, name))
+func (s *LaunchTemplateServiceImpl) getHeaderByName(ctx context.Context, accountID, name string) (*LaunchTemplateHeader, jetstream.KeyValueEntry, error) {
+	entry, err := s.kv.Get(ctx, nameKey(accountID, name))
 	if err != nil {
 		return nil, nil, errors.New(awserrors.ErrorInvalidLaunchTemplateNameNotFoundException)
 	}
 	ltID := string(entry.Value())
-	h, hentry, err := s.getHeaderByID(accountID, ltID)
+	h, hentry, err := s.getHeaderByID(ctx, accountID, ltID)
 	if err != nil {
 		return nil, nil, errors.New(awserrors.ErrorInvalidLaunchTemplateNameNotFoundException)
 	}
@@ -131,7 +133,7 @@ func (s *LaunchTemplateServiceImpl) getHeaderByName(accountID, name string) (*La
 
 // resolveHeader resolves a template by id or name (mutually exclusive) and
 // returns the header plus its KV entry for CAS operations.
-func (s *LaunchTemplateServiceImpl) resolveHeader(accountID string, ltID, name *string) (*LaunchTemplateHeader, nats.KeyValueEntry, error) {
+func (s *LaunchTemplateServiceImpl) resolveHeader(ctx context.Context, accountID string, ltID, name *string) (*LaunchTemplateHeader, jetstream.KeyValueEntry, error) {
 	id := aws.StringValue(ltID)
 	nm := aws.StringValue(name)
 	switch {
@@ -141,9 +143,9 @@ func (s *LaunchTemplateServiceImpl) resolveHeader(accountID string, ltID, name *
 		if err := validateTemplateID(id); err != nil {
 			return nil, nil, err
 		}
-		return s.getHeaderByID(accountID, id)
+		return s.getHeaderByID(ctx, accountID, id)
 	case nm != "":
-		return s.getHeaderByName(accountID, nm)
+		return s.getHeaderByName(ctx, accountID, nm)
 	default:
 		return nil, nil, errors.New(awserrors.ErrorMissingParameter)
 	}
@@ -151,8 +153,8 @@ func (s *LaunchTemplateServiceImpl) resolveHeader(accountID string, ltID, name *
 
 // getVersion loads an immutable version body. A missing body is the read-side
 // guard that turns a dangling default or deleted version into VersionNotFound.
-func (s *LaunchTemplateServiceImpl) getVersion(accountID, ltID string, n int64) (*LaunchTemplateVersionRec, error) {
-	entry, err := s.kv.Get(versionKey(accountID, ltID, n))
+func (s *LaunchTemplateServiceImpl) getVersion(ctx context.Context, accountID, ltID string, n int64) (*LaunchTemplateVersionRec, error) {
+	entry, err := s.kv.Get(ctx, versionKey(accountID, ltID, n))
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidLaunchTemplateIdVersionNotFound)
 	}
@@ -165,10 +167,10 @@ func (s *LaunchTemplateServiceImpl) getVersion(accountID, ltID string, n int64)
 
 // listVersionNumbers returns the existing version numbers for a template,
 // derived by scanning its version keys (sorted ascending).
-func (s *LaunchTemplateServiceImpl) listVersionNumbers(accountID, ltID string) ([]int64, error) {
-	keys, err := s.kv.Keys()
+func (s *LaunchTemplateServiceImpl) listVersionNumbers(ctx context.Context, accountID, ltID string) ([]int64, error) {
+	keys, err := s.kv.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, nil
 		}
 		return nil, errors.New(awserrors.ErrorServerInternal)
@@ -190,8 +192,8 @@ func (s *LaunchTemplateServiceImpl) listVersionNumbers(accountID, ltID string) (
 }
 
 // latestVersionNumber returns the highest existing version number, or 0 if none.
-func (s *LaunchTemplateServiceImpl) latestVersionNumber(accountID, ltID string) (int64, error) {
-	nums, err := s.listVersionNumbers(accountID, ltID)
+func (s *LaunchTemplateServiceImpl) latestVersionNumber(ctx context.Context, accountID, ltID string) (int64, error) {
+	nums, err := s.listVersionNumbers(ctx, accountID, ltID)
 	if err != nil {
 		return 0, err
 	}
@@ -235,12 +237,12 @@ func latestVersionsFromKeys(keys []string, prefix string) map[string]int64 {
 // resolveVersionNumber maps a selector ("", $Default, $Latest, or numeric) to a
 // concrete version number. It does not verify the body exists — callers load the
 // body and surface VersionNotFound when it is missing.
-func (s *LaunchTemplateServiceImpl) resolveVersionNumber(accountID string, h *LaunchTemplateHeader, sel string) (int64, error) {
+func (s *LaunchTemplateServiceImpl) resolveVersionNumber(ctx context.Context, accountID string, h *LaunchTemplateHeader, sel string) (int64, error) {
 	switch sel {
 	case "", versionDefault:
 		return h.DefaultVersionNumber, nil
 	case versionLatest:
-		return s.latestVersionNumber(accountID, h.LaunchTemplateId)
+		return s.latestVersionNumber(ctx, accountID, h.LaunchTemplateId)
 	default:
 		n, err := strconv.ParseInt(sel, 10, 64)
 		if err != nil {
@@ -270,7 +272,7 @@ func (s *LaunchTemplateServiceImpl) CreateLaunchTemplate(ctx context.Context, in
 	}
 
 	ltID := utils.GenerateResourceID("lt")
-	if err := s.claimName(accountID, name, ltID); err != nil {
+	if err := s.claimName(ctx, accountID, name, ltID); err != nil {
 		return nil, err
 	}
 
@@ -282,7 +284,7 @@ func (s *LaunchTemplateServiceImpl) CreateLaunchTemplate(ctx context.Context, in
 		CreatedBy:          accountID,
 		Data:               data,
 	}
-	if err := s.putVersion(accountID, ltID, &rec); err != nil {
+	if err := s.putVersion(ctx, accountID, ltID, &rec); err != nil {
 		return nil, err
 	}
 
@@ -295,7 +297,7 @@ func (s *LaunchTemplateServiceImpl) CreateLaunchTemplate(ctx context.Context, in
 		DefaultVersionNumber: 1,
 		Tags:                 utils.ExtractTags(input.TagSpecifications, launchTemplateTagResourceType),
 	}
-	if err := s.putHeader(accountID, &header); err != nil {
+	if err := s.putHeader(ctx, accountID, &header); err != nil {
 		return nil, err
 	}
 
@@ -310,12 +312,12 @@ func (s *LaunchTemplateServiceImpl) CreateLaunchTemplate(ctx context.Context, in
 // repair-on-write: if the name already exists but its header is gone (a crash
 // orphan), reclaim it via a revision-guarded CAS update. Only a name whose header
 // still exists returns AlreadyExists.
-func (s *LaunchTemplateServiceImpl) claimName(accountID, name, ltID string) error {
+func (s *LaunchTemplateServiceImpl) claimName(ctx context.Context, accountID, name, ltID string) error {
 	key := nameKey(accountID, name)
-	if _, err := s.kv.Create(key, []byte(ltID)); err == nil {
+	if _, err := s.kv.Create(ctx, key, []byte(ltID)); err == nil {
 		return nil
 	}
-	entry, err := s.kv.Get(key)
+	entry, err := s.kv.Get(ctx, key)
 	if err != nil {
 		return errors.New(awserrors.ErrorServerInternal)
 	}
@@ -323,36 +325,36 @@ func (s *LaunchTemplateServiceImpl) claimName(accountID, name, ltID string) erro
 	// header means the name is taken; any other read error (transient fault, or a
 	// concurrent in-flight create whose header is not yet written) fails closed so
 	// the name is never stolen from a possibly-live template.
-	_, herr := s.kv.Get(headerKey(accountID, string(entry.Value())))
+	_, herr := s.kv.Get(ctx, headerKey(accountID, string(entry.Value())))
 	switch {
 	case herr == nil:
 		return errors.New(awserrors.ErrorInvalidLaunchTemplateNameAlreadyExistsException)
-	case !errors.Is(herr, nats.ErrKeyNotFound):
+	case !errors.Is(herr, jetstream.ErrKeyNotFound):
 		return errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.kv.Update(key, []byte(ltID), entry.Revision()); err != nil {
+	if _, err := s.kv.Update(ctx, key, []byte(ltID), entry.Revision()); err != nil {
 		return errors.New(awserrors.ErrorInvalidLaunchTemplateNameAlreadyExistsException)
 	}
 	return nil
 }
 
-func (s *LaunchTemplateServiceImpl) putHeader(accountID string, h *LaunchTemplateHeader) error {
+func (s *LaunchTemplateServiceImpl) putHeader(ctx context.Context, accountID string, h *LaunchTemplateHeader) error {
 	data, err := json.Marshal(h)
 	if err != nil {
 		return errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.kv.Put(headerKey(accountID, h.LaunchTemplateId), data); err != nil {
+	if _, err := s.kv.Put(ctx, headerKey(accountID, h.LaunchTemplateId), data); err != nil {
 		return errors.New(awserrors.ErrorServerInternal)
 	}
 	return nil
 }
 
-func (s *LaunchTemplateServiceImpl) putVersion(accountID, ltID string, rec *LaunchTemplateVersionRec) error {
+func (s *LaunchTemplateServiceImpl) putVersion(ctx context.Context, accountID, ltID string, rec *LaunchTemplateVersionRec) error {
 	data, err := json.Marshal(rec)
 	if err != nil {
 		return errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.kv.Put(versionKey(accountID, ltID, rec.VersionNumber), data); err != nil {
+	if _, err := s.kv.Put(ctx, versionKey(accountID, ltID, rec.VersionNumber), data); err != nil {
 		return errors.New(awserrors.ErrorServerInternal)
 	}
 	return nil
@@ -364,7 +366,7 @@ func (s *LaunchTemplateServiceImpl) CreateLaunchTemplateVersion(ctx context.Cont
 	if input.LaunchTemplateData == nil {
 		return nil, errors.New(awserrors.ErrorMissingParameter)
 	}
-	header, _, err := s.resolveHeader(accountID, input.LaunchTemplateId, input.LaunchTemplateName)
+	header, _, err := s.resolveHeader(ctx, accountID, input.LaunchTemplateId, input.LaunchTemplateName)
 	if err != nil {
 		return nil, err
 	}
@@ -380,11 +382,11 @@ func (s *LaunchTemplateServiceImpl) CreateLaunchTemplateVersion(ctx context.Cont
 	// SourceVersion clones-and-overrides; without it the new version is the data as given.
 	data := override
 	if src := aws.StringValue(input.SourceVersion); src != "" {
-		n, err := s.resolveVersionNumber(accountID, header, src)
+		n, err := s.resolveVersionNumber(ctx, accountID, header, src)
 		if err != nil {
 			return nil, err
 		}
-		base, err := s.getVersion(accountID, header.LaunchTemplateId, n)
+		base, err := s.getVersion(ctx, accountID, header.LaunchTemplateId, n)
 		if err != nil {
 			return nil, err
 		}
@@ -401,7 +403,7 @@ func (s *LaunchTemplateServiceImpl) CreateLaunchTemplateVersion(ctx context.Cont
 	// Assign the number with an atomic kv.Create in a bounded retry loop: on
 	// collision another writer took n, so rescan and try the next number.
 	for attempt := range maxVersionRetries {
-		latest, err := s.latestVersionNumber(accountID, header.LaunchTemplateId)
+		latest, err := s.latestVersionNumber(ctx, accountID, header.LaunchTemplateId)
 		if err != nil {
 			return nil, err
 		}
@@ -411,8 +413,8 @@ func (s *LaunchTemplateServiceImpl) CreateLaunchTemplateVersion(ctx context.Cont
 		if err != nil {
 			return nil, errors.New(awserrors.ErrorServerInternal)
 		}
-		if _, err := s.kv.Create(versionKey(accountID, header.LaunchTemplateId, next), body); err != nil {
-			if errors.Is(err, nats.ErrKeyExists) {
+		if _, err := s.kv.Create(ctx, versionKey(accountID, header.LaunchTemplateId, next), body); err != nil {
+			if errors.Is(err, jetstream.ErrKeyExists) {
 				slog.DebugContext(ctx, "CreateLaunchTemplateVersion: version-number collision, retrying", "attempt", attempt, "next", next)
 				continue
 			}
@@ -432,12 +434,12 @@ func (s *LaunchTemplateServiceImpl) CreateLaunchTemplateVersion(ctx context.Cont
 // --- ModifyLaunchTemplate ---
 
 func (s *LaunchTemplateServiceImpl) ModifyLaunchTemplate(ctx context.Context, input *ec2.ModifyLaunchTemplateInput, accountID string) (*ec2.ModifyLaunchTemplateOutput, error) {
-	header, entry, err := s.resolveHeader(accountID, input.LaunchTemplateId, input.LaunchTemplateName)
+	header, entry, err := s.resolveHeader(ctx, accountID, input.LaunchTemplateId, input.LaunchTemplateName)
 	if err != nil {
 		return nil, err
 	}
 	if aws.BoolValue(input.DryRun) {
-		latest, err := s.latestVersionNumber(accountID, header.LaunchTemplateId)
+		latest, err := s.latestVersionNumber(ctx, accountID, header.LaunchTemplateId)
 		if err != nil {
 			return nil, err
 		}
@@ -445,13 +447,13 @@ func (s *LaunchTemplateServiceImpl) ModifyLaunchTemplate(ctx context.Context, in
 	}
 
 	if sel := aws.StringValue(input.DefaultVersion); sel != "" {
-		n, err := s.resolveVersionNumber(accountID, header, sel)
+		n, err := s.resolveVersionNumber(ctx, accountID, header, sel)
 		if err != nil {
 			return nil, err
 		}
 		// Verify the target body exists immediately before the header CAS so a
 		// concurrent delete degrades to VersionNotFound, never a dangling default.
-		if _, err := s.getVersion(accountID, header.LaunchTemplateId, n); err != nil {
+		if _, err := s.getVersion(ctx, accountID, header.LaunchTemplateId, n); err != nil {
 			return nil, err
 		}
 		header.DefaultVersionNumber = n
@@ -459,12 +461,12 @@ func (s *LaunchTemplateServiceImpl) ModifyLaunchTemplate(ctx context.Context, in
 		if err != nil {
 			return nil, errors.New(awserrors.ErrorServerInternal)
 		}
-		if _, err := s.kv.Update(headerKey(accountID, header.LaunchTemplateId), data, entry.Revision()); err != nil {
+		if _, err := s.kv.Update(ctx, headerKey(accountID, header.LaunchTemplateId), data, entry.Revision()); err != nil {
 			return nil, errors.New(awserrors.ErrorServerInternal)
 		}
 	}
 
-	latest, err := s.latestVersionNumber(accountID, header.LaunchTemplateId)
+	latest, err := s.latestVersionNumber(ctx, accountID, header.LaunchTemplateId)
 	if err != nil {
 		return nil, err
 	}
@@ -475,19 +477,19 @@ func (s *LaunchTemplateServiceImpl) ModifyLaunchTemplate(ctx context.Context, in
 // --- DeleteLaunchTemplate ---
 
 func (s *LaunchTemplateServiceImpl) DeleteLaunchTemplate(ctx context.Context, input *ec2.DeleteLaunchTemplateInput, accountID string) (*ec2.DeleteLaunchTemplateOutput, error) {
-	header, _, err := s.resolveHeader(accountID, input.LaunchTemplateId, input.LaunchTemplateName)
+	header, _, err := s.resolveHeader(ctx, accountID, input.LaunchTemplateId, input.LaunchTemplateName)
 	if err != nil {
 		return nil, err
 	}
 	if aws.BoolValue(input.DryRun) {
-		latest, err := s.latestVersionNumber(accountID, header.LaunchTemplateId)
+		latest, err := s.latestVersionNumber(ctx, accountID, header.LaunchTemplateId)
 		if err != nil {
 			return nil, err
 		}
 		return &ec2.DeleteLaunchTemplateOutput{LaunchTemplate: headerToEC2(header, latest)}, nil
 	}
 
-	latest, err := s.latestVersionNumber(accountID, header.LaunchTemplateId)
+	latest, err := s.latestVersionNumber(ctx, accountID, header.LaunchTemplateId)
 	if err != nil {
 		return nil, err
 	}
@@ -495,18 +497,18 @@ func (s *LaunchTemplateServiceImpl) DeleteLaunchTemplate(ctx context.Context, in
 
 	// Delete the header first: the template immediately vanishes from every
 	// describe. Version bodies and the name index are best-effort cleanup.
-	if err := s.kv.Delete(headerKey(accountID, header.LaunchTemplateId)); err != nil {
+	if err := s.kv.Delete(ctx, headerKey(accountID, header.LaunchTemplateId)); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
-	if err := s.kv.Delete(nameKey(accountID, header.LaunchTemplateName)); err != nil {
+	if err := s.kv.Delete(ctx, nameKey(accountID, header.LaunchTemplateName)); err != nil {
 		slog.WarnContext(ctx, "DeleteLaunchTemplate: name index cleanup failed", "name", header.LaunchTemplateName, "err", err)
 	}
-	nums, err := s.listVersionNumbers(accountID, header.LaunchTemplateId)
+	nums, err := s.listVersionNumbers(ctx, accountID, header.LaunchTemplateId)
 	if err != nil {
 		slog.WarnContext(ctx, "DeleteLaunchTemplate: version enumeration failed, bodies left in place", "launchTemplateId", header.LaunchTemplateId, "err", err)
 	}
 	for _, n := range nums {
-		if err := s.kv.Delete(versionKey(accountID, header.LaunchTemplateId, n)); err != nil {
+		if err := s.kv.Delete(ctx, versionKey(accountID, header.LaunchTemplateId, n)); err != nil {
 			slog.WarnContext(ctx, "DeleteLaunchTemplate: version cleanup failed", "launchTemplateId", header.LaunchTemplateId, "version", n, "err", err)
 		}
 	}
@@ -518,7 +520,7 @@ func (s *LaunchTemplateServiceImpl) DeleteLaunchTemplate(ctx context.Context, in
 // --- DeleteLaunchTemplateVersions ---
 
 func (s *LaunchTemplateServiceImpl) DeleteLaunchTemplateVersions(ctx context.Context, input *ec2.DeleteLaunchTemplateVersionsInput, accountID string) (*ec2.DeleteLaunchTemplateVersionsOutput, error) {
-	header, _, err := s.resolveHeader(accountID, input.LaunchTemplateId, input.LaunchTemplateName)
+	header, _, err := s.resolveHeader(ctx, accountID, input.LaunchTemplateId, input.LaunchTemplateName)
 	if err != nil {
 		return nil, err
 	}
@@ -544,12 +546,12 @@ func (s *LaunchTemplateServiceImpl) DeleteLaunchTemplateVersions(ctx context.Con
 				deleteErrorItem(header, n, awserrors.ErrorInvalidParameterValue, "cannot delete the default version of a launch template"))
 			continue
 		}
-		if _, err := s.getVersion(accountID, header.LaunchTemplateId, n); err != nil {
+		if _, err := s.getVersion(ctx, accountID, header.LaunchTemplateId, n); err != nil {
 			out.UnsuccessfullyDeletedLaunchTemplateVersions = append(out.UnsuccessfullyDeletedLaunchTemplateVersions,
 				deleteErrorItem(header, n, awserrors.ErrorInvalidLaunchTemplateIdVersionNotFound, "version does not exist"))
 			continue
 		}
-		if err := s.kv.Delete(versionKey(accountID, header.LaunchTemplateId, n)); err != nil {
+		if err := s.kv.Delete(ctx, versionKey(accountID, header.LaunchTemplateId, n)); err != nil {
 			out.UnsuccessfullyDeletedLaunchTemplateVersions = append(out.UnsuccessfullyDeletedLaunchTemplateVersions,
 				deleteErrorItem(header, n, awserrors.ErrorServerInternal, "failed to delete version"))
 			continue
@@ -604,8 +606,8 @@ func (s *LaunchTemplateServiceImpl) DescribeLaunchTemplates(ctx context.Context,
 	idSet := stringSet(input.LaunchTemplateIds)
 
 	prefix := accountID + "."
-	keys, err := s.kv.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	keys, err := s.kv.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -620,7 +622,7 @@ func (s *LaunchTemplateServiceImpl) DescribeLaunchTemplates(ctx context.Context,
 		if !isHeaderKey(k, prefix) {
 			continue
 		}
-		entry, err := s.kv.Get(k)
+		entry, err := s.kv.Get(ctx, k)
 		if err != nil {
 			slog.WarnContext(ctx, "DescribeLaunchTemplates: header read failed", "key", k, "err", err)
 			continue
@@ -713,19 +715,19 @@ func (s *LaunchTemplateServiceImpl) DescribeLaunchTemplateVersions(ctx context.C
 		return nil, errors.New(awserrors.ErrorInvalidParameterValue)
 	}
 
-	header, _, err := s.resolveHeader(accountID, input.LaunchTemplateId, input.LaunchTemplateName)
+	header, _, err := s.resolveHeader(ctx, accountID, input.LaunchTemplateId, input.LaunchTemplateName)
 	if err != nil {
 		return nil, err
 	}
 
-	numbers, err := s.selectVersionNumbers(accountID, header, input)
+	numbers, err := s.selectVersionNumbers(ctx, accountID, header, input)
 	if err != nil {
 		return nil, err
 	}
 
 	var versions []*ec2.LaunchTemplateVersion
 	for _, n := range numbers {
-		rec, err := s.getVersion(accountID, header.LaunchTemplateId, n)
+		rec, err := s.getVersion(ctx, accountID, header.LaunchTemplateId, n)
 		if err != nil {
 			return nil, err
 		}
@@ -742,16 +744,16 @@ func (s *LaunchTemplateServiceImpl) DescribeLaunchTemplateVersions(ctx context.C
 // selectVersionNumbers resolves the requested Versions/MinVersion/MaxVersion into
 // a concrete, ordered list of version numbers. Explicit Versions each resolve and
 // must exist; otherwise all versions within the optional [Min,Max] range.
-func (s *LaunchTemplateServiceImpl) selectVersionNumbers(accountID string, header *LaunchTemplateHeader, input *ec2.DescribeLaunchTemplateVersionsInput) ([]int64, error) {
+func (s *LaunchTemplateServiceImpl) selectVersionNumbers(ctx context.Context, accountID string, header *LaunchTemplateHeader, input *ec2.DescribeLaunchTemplateVersionsInput) ([]int64, error) {
 	if len(input.Versions) > 0 {
 		seen := make(map[int64]bool)
 		var out []int64
 		for _, v := range input.Versions {
-			n, err := s.resolveVersionNumber(accountID, header, aws.StringValue(v))
+			n, err := s.resolveVersionNumber(ctx, accountID, header, aws.StringValue(v))
 			if err != nil {
 				return nil, err
 			}
-			if _, err := s.getVersion(accountID, header.LaunchTemplateId, n); err != nil {
+			if _, err := s.getVersion(ctx, accountID, header.LaunchTemplateId, n); err != nil {
 				return nil, err
 			}
 			if !seen[n] {
@@ -762,7 +764,7 @@ func (s *LaunchTemplateServiceImpl) selectVersionNumbers(accountID string, heade
 		return out, nil
 	}
 
-	all, err := s.listVersionNumbers(accountID, header.LaunchTemplateId)
+	all, err := s.listVersionNumbers(ctx, accountID, header.LaunchTemplateId)
 	if err != nil {
 		return nil, err
 	}

spinifex/handlers/ec2/launchtemplate/service_impl_test.go

@@ -21,7 +21,7 @@ const (
 func setupTestService(t *testing.T) *LaunchTemplateServiceImpl {
 	t.Helper()
 	_, nc, _ := testutil.StartTestJetStream(t)
-	svc, err := NewLaunchTemplateServiceImplWithNATS(nil, nc)
+	svc, err := NewLaunchTemplateServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 	return svc
 }
@@ -129,7 +129,7 @@ func TestCreateLaunchTemplate_OrphanNameReclaim(t *testing.T) {
 	lt := createTemplate(t, svc, "orphan", "t3.micro")
 
 	// Simulate a crash after the name claim but with the header lost.
-	require.NoError(t, svc.kv.Delete(headerKey(testAccountID, aws.StringValue(lt.LaunchTemplateId))))
+	require.NoError(t, svc.kv.Delete(t.Context(), headerKey(testAccountID, aws.StringValue(lt.LaunchTemplateId))))
 
 	// The name index still points at the now-orphaned id; create must reclaim it.
 	out, err := svc.CreateLaunchTemplate(context.Background(), &ec2.CreateLaunchTemplateInput{
@@ -459,10 +459,10 @@ func TestDeleteLaunchTemplate_RemovesEverything(t *testing.T) {
 	assert.Equal(t, awserrors.ErrorInvalidLaunchTemplateIdNotFound, err.Error())
 
 	// Version bodies and name index gone.
-	nums, err := svc.listVersionNumbers(testAccountID, id)
+	nums, err := svc.listVersionNumbers(t.Context(), testAccountID, id)
 	require.NoError(t, err)
 	assert.Empty(t, nums)
-	_, err = svc.kv.Get(nameKey(testAccountID, "web"))
+	_, err = svc.kv.Get(t.Context(), nameKey(testAccountID, "web"))
 	require.Error(t, err)
 
 	// Name is reusable after delete.

spinifex/handlers/ec2/natgw/lifecycle_invariants_test.go

@@ -1,7 +1,6 @@
 package handlers_ec2_natgw
 
 import (
-	"context"
 	"encoding/json"
 	"fmt"
 	"testing"
@@ -10,7 +9,6 @@ import (
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	handlers_ec2_eip "github.com/mulgadc/spinifex/spinifex/handlers/ec2/eip"
-	"github.com/mulgadc/spinifex/spinifex/testutil"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -24,7 +22,7 @@ import (
 func TestRLC1_NatGatewayDeleteNotFoundOnAbsent(t *testing.T) {
 	svc := setupTestService(t)
 
-	_, err := svc.DeleteNatGateway(context.Background(), &ec2.DeleteNatGatewayInput{
+	_, err := svc.DeleteNatGateway(t.Context(), &ec2.DeleteNatGatewayInput{
 		NatGatewayId: aws.String("nat-absent00000000"),
 	}, testAccountID)
 
@@ -42,12 +40,12 @@ func TestRLC3_NatGatewayDeletesWhileRouted(t *testing.T) {
 		svc, js := setupTestServiceJS(t)
 		natgwID := createTestNatGateway(t, svc)
 
-		testutil.SeedKV(t, js, kvBucketRouteTables, map[string][]byte{
+		seedKV(t, js, kvBucketRouteTables, map[string][]byte{
 			utils.AccountKey(testAccountID, "rtb-routed"): fmt.Appendf(nil,
 				`{"route_table_id":"rtb-routed","vpc_id":"vpc-test1","routes":[{"destination_cidr_block":"0.0.0.0/0","nat_gateway_id":%q}]}`, natgwID),
 		})
 
-		_, err := svc.DeleteNatGateway(context.Background(), &ec2.DeleteNatGatewayInput{NatGatewayId: aws.String(natgwID)}, testAccountID)
+		_, err := svc.DeleteNatGateway(t.Context(), &ec2.DeleteNatGatewayInput{NatGatewayId: aws.String(natgwID)}, testAccountID)
 		require.NoErrorf(t, err,
 			"ADR-0004 §2: DeleteNatGateway must succeed even while a live route table forwards to it (the route blackholes, AWS-faithful)")
 	})
@@ -56,7 +54,7 @@ func TestRLC3_NatGatewayDeletesWhileRouted(t *testing.T) {
 		svc := setupTestService(t)
 		natgwID := createTestNatGateway(t, svc)
 
-		_, err := svc.DeleteNatGateway(context.Background(), &ec2.DeleteNatGatewayInput{NatGatewayId: aws.String(natgwID)}, testAccountID)
+		_, err := svc.DeleteNatGateway(t.Context(), &ec2.DeleteNatGatewayInput{NatGatewayId: aws.String(natgwID)}, testAccountID)
 		require.NoErrorf(t, err,
 			"ADR-0004 §2: with no route forwarding to it, DeleteNatGateway must succeed")
 	})
@@ -74,7 +72,7 @@ func TestRLC3_NatGatewayDeleteDisassociatesEIPKeepingAllocated(t *testing.T) {
 	require.Equal(t, "associated", assoc.State, "CreateNatGateway must mark the EIP associated so a second NAT GW cannot reuse it")
 	require.NotEmpty(t, assoc.AssociationId)
 
-	_, err := svc.DeleteNatGateway(context.Background(), &ec2.DeleteNatGatewayInput{NatGatewayId: aws.String(natgwID)}, testAccountID)
+	_, err := svc.DeleteNatGateway(t.Context(), &ec2.DeleteNatGatewayInput{NatGatewayId: aws.String(natgwID)}, testAccountID)
 	require.NoError(t, err)
 
 	eip := readEIP(t, svc, "eipalloc-test1")
@@ -84,7 +82,7 @@ func TestRLC3_NatGatewayDeleteDisassociatesEIPKeepingAllocated(t *testing.T) {
 
 func createTestNatGateway(t *testing.T, svc *NatGatewayServiceImpl) string {
 	t.Helper()
-	out, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	out, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-test1"),
 	}, testAccountID)
@@ -94,7 +92,7 @@ func createTestNatGateway(t *testing.T, svc *NatGatewayServiceImpl) string {
 
 func readEIP(t *testing.T, svc *NatGatewayServiceImpl, allocID string) handlers_ec2_eip.EIPRecord {
 	t.Helper()
-	entry, err := svc.eipKV.Get(utils.AccountKey(testAccountID, allocID))
+	entry, err := svc.eipKV.Get(t.Context(), utils.AccountKey(testAccountID, allocID))
 	require.NoError(t, err)
 	var eip handlers_ec2_eip.EIPRecord
 	require.NoError(t, json.Unmarshal(entry.Value(), &eip))

spinifex/handlers/ec2/natgw/record_tags_test.go

@@ -1,7 +1,6 @@
 package handlers_ec2_natgw
 
 import (
-	"context"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -31,7 +30,7 @@ func TestNatGatewayRecordTagsMirror(t *testing.T) {
 		},
 	}, testAccountID))
 
-	out, err := svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err := svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		NatGatewayIds: []*string{aws.String(natgwID)},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -47,7 +46,7 @@ func TestNatGatewayRecordTagsMirror(t *testing.T) {
 		},
 	}, testAccountID))
 
-	out, err = svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err = svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		NatGatewayIds: []*string{aws.String(natgwID)},
 	}, testAccountID)
 	require.NoError(t, err)

spinifex/handlers/ec2/natgw/service_impl.go

@@ -15,67 +15,59 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
 	handlers_ec2_eip "github.com/mulgadc/spinifex/spinifex/handlers/ec2/eip"
 	handlers_ec2_vpc "github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 var _ NatGatewayService = (*NatGatewayServiceImpl)(nil)
 
 // NatGatewayServiceImpl implements NAT Gateway operations with NATS JetStream persistence.
 type NatGatewayServiceImpl struct {
-	natgwKV        nats.KeyValue
-	deletedNatgwKV nats.KeyValue
-	eipKV          nats.KeyValue
-	subnetKV       nats.KeyValue
-	vpcKV          nats.KeyValue
+	natgwKV        jetstream.KeyValue
+	deletedNatgwKV jetstream.KeyValue
+	eipKV          jetstream.KeyValue
+	subnetKV       jetstream.KeyValue
+	vpcKV          jetstream.KeyValue
 	natsConn       *nats.Conn
 }
 
 // NewNatGatewayServiceImplWithNATS creates a NAT Gateway service.
-func NewNatGatewayServiceImplWithNATS(natsConn *nats.Conn) (*NatGatewayServiceImpl, error) {
-	js, err := natsConn.JetStream()
+func NewNatGatewayServiceImplWithNATS(ctx context.Context, natsConn *nats.Conn) (*NatGatewayServiceImpl, error) {
+	js, err := jetstream.New(natsConn)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get JetStream context: %w", err)
 	}
 
-	natgwKV, err := utils.GetOrCreateKVBucket(js, KVBucketNatGateways, 10)
+	natgwKV, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketNatGateways, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketNatGateways, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketNatGateways, natgwKV, KVBucketNatGatewaysVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketNatGateways, natgwKV, KVBucketNatGatewaysVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketNatGateways, err)
 	}
 
-	// Deleted NAT Gateways bucket with 1-hour TTL — keys auto-expire.
-	// Terraform polls DescribeNatGateways after delete and expects state=deleted.
-	deletedKV, err := js.CreateKeyValue(&nats.KeyValueConfig{
-		Bucket:      KVBucketDeletedNatGateways,
-		Description: "Deleted NAT Gateways (auto-expire after 1 hour)",
-		History:     1,
-		TTL:         1 * time.Hour,
-	})
+	deletedKV, err := getOrCreateDeletedBucket(ctx, js)
 	if err != nil {
-		deletedKV, err = js.KeyValue(KVBucketDeletedNatGateways)
-		if err != nil {
-			return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketDeletedNatGateways, err)
-		}
+		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketDeletedNatGateways, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketDeletedNatGateways, deletedKV, KVBucketDeletedNatGatewaysVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketDeletedNatGateways, deletedKV, KVBucketDeletedNatGatewaysVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketDeletedNatGateways, err)
 	}
 
-	eipKV, err := utils.GetOrCreateKVBucket(js, handlers_ec2_eip.KVBucketEIPs, 10)
+	eipKV, err := kvutil.GetOrCreateBucket(ctx, js, handlers_ec2_eip.KVBucketEIPs, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get EIP KV bucket: %w", err)
 	}
 
-	subnetKV, err := utils.GetOrCreateKVBucket(js, handlers_ec2_vpc.KVBucketSubnets, 10)
+	subnetKV, err := kvutil.GetOrCreateBucket(ctx, js, handlers_ec2_vpc.KVBucketSubnets, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get subnet KV bucket: %w", err)
 	}
 
-	vpcKV, err := utils.GetOrCreateKVBucket(js, handlers_ec2_vpc.KVBucketVPCs, 10)
+	vpcKV, err := kvutil.GetOrCreateBucket(ctx, js, handlers_ec2_vpc.KVBucketVPCs, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get VPC KV bucket: %w", err)
 	}
@@ -92,6 +84,25 @@ func NewNatGatewayServiceImplWithNATS(natsConn *nats.Conn) (*NatGatewayServiceIm
 	}, nil
 }
 
+// getOrCreateDeletedBucket creates the expiring bucket without changing an
+// existing bucket's replica or placement configuration.
+func getOrCreateDeletedBucket(ctx context.Context, js jetstream.KeyValueManager) (jetstream.KeyValue, error) {
+	// Terraform polls DescribeNatGateways after delete and expects state=deleted.
+	kv, err := js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
+		Bucket:      KVBucketDeletedNatGateways,
+		Description: "Deleted NAT Gateways (auto-expire after 1 hour)",
+		History:     1,
+		TTL:         time.Hour,
+	})
+	if err == nil {
+		return kv, nil
+	}
+	if !errors.Is(err, jetstream.ErrBucketExists) {
+		return nil, err
+	}
+	return js.KeyValue(ctx, KVBucketDeletedNatGateways)
+}
+
 // natGatewayEvent is published on vpc.add-nat-gateway / vpc.delete-nat-gateway topics.
 type natGatewayEvent struct {
 	VpcId        string `json:"vpc_id"`
@@ -113,7 +124,7 @@ func (s *NatGatewayServiceImpl) CreateNatGateway(ctx context.Context, input *ec2
 	allocID := *input.AllocationId
 
 	// Validate subnet exists and get its VPC
-	subnetEntry, err := s.subnetKV.Get(utils.AccountKey(accountID, subnetID))
+	subnetEntry, err := s.subnetKV.Get(ctx, utils.AccountKey(accountID, subnetID))
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidSubnetIDNotFound)
 	}
@@ -123,7 +134,7 @@ func (s *NatGatewayServiceImpl) CreateNatGateway(ctx context.Context, input *ec2
 	}
 
 	// Validate EIP exists and is not already associated
-	eipEntry, err := s.eipKV.Get(utils.AccountKey(accountID, allocID))
+	eipEntry, err := s.eipKV.Get(ctx, utils.AccountKey(accountID, allocID))
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidAllocationIDNotFound)
 	}
@@ -152,7 +163,7 @@ func (s *NatGatewayServiceImpl) CreateNatGateway(ctx context.Context, input *ec2
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.natgwKV.Put(utils.AccountKey(accountID, natgwID), data); err != nil {
+	if _, err := s.natgwKV.Put(ctx, utils.AccountKey(accountID, natgwID), data); err != nil {
 		slog.ErrorContext(ctx, "Failed to store NAT Gateway", "natGatewayId", natgwID, "err", err)
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
@@ -165,7 +176,7 @@ func (s *NatGatewayServiceImpl) CreateNatGateway(ctx context.Context, input *ec2
 	eipRecord.VpcId = subnetRecord.VpcId
 	eipRecord.State = "associated"
 	if eipData, err := json.Marshal(eipRecord); err == nil {
-		if _, err := s.eipKV.Update(utils.AccountKey(accountID, allocID), eipData, eipEntry.Revision()); err != nil {
+		if _, err := s.eipKV.Update(ctx, utils.AccountKey(accountID, allocID), eipData, eipEntry.Revision()); err != nil {
 			slog.WarnContext(ctx, "CreateNatGateway: failed to mark EIP associated", "natGatewayId", natgwID, "allocationId", allocID, "err", err)
 		}
 	}
@@ -187,9 +198,9 @@ func (s *NatGatewayServiceImpl) DeleteNatGateway(ctx context.Context, input *ec2
 	natgwID := *input.NatGatewayId
 	key := utils.AccountKey(accountID, natgwID)
 
-	entry, err := s.natgwKV.Get(key)
+	entry, err := s.natgwKV.Get(ctx, key)
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			// AWS-faithful: an absent NAT gateway is NotFound (provider
 			// tolerates it on destroy); destroy orchestration tolerates it too.
 			return nil, errors.New(awserrors.ErrorInvalidNatGatewayIDNotFound)
@@ -215,13 +226,13 @@ func (s *NatGatewayServiceImpl) DeleteNatGateway(ctx context.Context, input *ec2
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.deletedNatgwKV.Put(key, deleted); err != nil {
+	if _, err := s.deletedNatgwKV.Put(ctx, key, deleted); err != nil {
 		slog.ErrorContext(ctx, "Failed to write deleted NAT Gateway", "natGatewayId", natgwID, "err", err)
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
 	// Remove from active bucket
-	if err := s.natgwKV.Delete(key); err != nil {
+	if err := s.natgwKV.Delete(ctx, key); err != nil {
 		slog.ErrorContext(ctx, "Failed to delete NAT Gateway from active bucket", "natGatewayId", natgwID, "err", err)
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
@@ -250,12 +261,17 @@ func (s *NatGatewayServiceImpl) publishDeleteEventsForNatGateway(ctx context.Con
 	if s.natsConn == nil {
 		return
 	}
-	rtbKV, err := utils.GetOrCreateKVBucket(mustJS(s.natsConn), kvBucketRouteTables, 10)
+	js, err := jetstream.New(s.natsConn)
+	if err != nil {
+		slog.WarnContext(ctx, "DeleteNatGateway: JetStream setup failed, SNAT teardown deferred to route delete", "natGatewayId", record.NatGatewayId, "err", err)
+		return
+	}
+	rtbKV, err := kvutil.GetOrCreateBucket(ctx, js, kvBucketRouteTables, 10)
 	if err != nil {
 		slog.WarnContext(ctx, "DeleteNatGateway: route-table scan failed, SNAT teardown deferred to route delete", "natGatewayId", record.NatGatewayId, "err", err)
 		return
 	}
-	keys, err := rtbKV.Keys()
+	keys, err := rtbKV.Keys(ctx)
 	if err != nil {
 		return
 	}
@@ -264,7 +280,7 @@ func (s *NatGatewayServiceImpl) publishDeleteEventsForNatGateway(ctx context.Con
 		if key == utils.VersionKey || !strings.HasPrefix(key, prefix) {
 			continue
 		}
-		entry, err := rtbKV.Get(key)
+		entry, err := rtbKV.Get(ctx, key)
 		if err != nil {
 			continue
 		}
@@ -297,7 +313,7 @@ func (s *NatGatewayServiceImpl) publishDeleteEventsForNatGateway(ctx context.Con
 			if assoc.SubnetId == "" {
 				continue
 			}
-			subnetEntry, err := s.subnetKV.Get(utils.AccountKey(accountID, assoc.SubnetId))
+			subnetEntry, err := s.subnetKV.Get(ctx, utils.AccountKey(accountID, assoc.SubnetId))
 			if err != nil {
 				continue
 			}
@@ -310,11 +326,6 @@ func (s *NatGatewayServiceImpl) publishDeleteEventsForNatGateway(ctx context.Con
 	}
 }
 
-func mustJS(nc *nats.Conn) nats.JetStreamContext {
-	js, _ := nc.JetStream()
-	return js
-}
-
 // disassociateEIP clears the EIP's association to this NAT gateway, keeping the
 // allocation (AWS parity). Idempotent and non-fatal: the NAT gateway is already
 // gone, so a stale association only delays the EIP's reuse.
@@ -323,9 +334,9 @@ func (s *NatGatewayServiceImpl) disassociateEIP(ctx context.Context, record *Nat
 		return
 	}
 	key := utils.AccountKey(accountID, record.AllocationId)
-	entry, err := s.eipKV.Get(key)
+	entry, err := s.eipKV.Get(ctx, key)
 	if err != nil {
-		if !errors.Is(err, nats.ErrKeyNotFound) {
+		if !errors.Is(err, jetstream.ErrKeyNotFound) {
 			slog.WarnContext(ctx, "DeleteNatGateway: EIP read failed during disassociate", "allocationId", record.AllocationId, "err", err)
 		}
 		return
@@ -348,7 +359,7 @@ func (s *NatGatewayServiceImpl) disassociateEIP(ctx context.Context, record *Nat
 	if err != nil {
 		return
 	}
-	if _, err := s.eipKV.Update(key, data, entry.Revision()); err != nil {
+	if _, err := s.eipKV.Update(ctx, key, data, entry.Revision()); err != nil {
 		slog.WarnContext(ctx, "DeleteNatGateway: failed to disassociate EIP", "allocationId", record.AllocationId, "err", err)
 	}
 }
@@ -376,8 +387,8 @@ func (s *NatGatewayServiceImpl) DescribeNatGateways(ctx context.Context, input *
 	}
 
 	prefix := accountID + "."
-	keys, err := s.natgwKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	keys, err := s.natgwKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -389,9 +400,9 @@ func (s *NatGatewayServiceImpl) DescribeNatGateways(ctx context.Context, input *
 			continue
 		}
 
-		entry, err := s.natgwKV.Get(key)
+		entry, err := s.natgwKV.Get(ctx, key)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
 			}
 			slog.ErrorContext(ctx, "Failed to read NAT Gateway", "key", key, "err", err)
@@ -421,9 +432,9 @@ func (s *NatGatewayServiceImpl) DescribeNatGateways(ctx context.Context, input *
 			continue
 		}
 		key := utils.AccountKey(accountID, id)
-		entry, err := s.deletedNatgwKV.Get(key)
+		entry, err := s.deletedNatgwKV.Get(ctx, key)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				return nil, errors.New(awserrors.ErrorInvalidNatGatewayIDNotFound)
 			}
 			return nil, errors.New(awserrors.ErrorServerInternal)
@@ -493,9 +504,10 @@ func (s *NatGatewayServiceImpl) PublishDeleteEvent(vpcId, natGatewayId, publicIp
 
 // GetNatGateway retrieves a NAT Gateway record by ID.
 func (s *NatGatewayServiceImpl) GetNatGateway(accountID, natgwID string) (*NatGatewayRecord, error) {
-	entry, err := s.natgwKV.Get(utils.AccountKey(accountID, natgwID))
+	ctx := context.Background()
+	entry, err := s.natgwKV.Get(ctx, utils.AccountKey(accountID, natgwID))
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return nil, errors.New(awserrors.ErrorInvalidNatGatewayIDNotFound)
 		}
 		return nil, errors.New(awserrors.ErrorServerInternal)
@@ -538,7 +550,7 @@ func (s *NatGatewayServiceImpl) ApplyRecordTags(input *ec2.CreateTagsInput, acco
 	if input == nil {
 		return nil
 	}
-	return utils.MirrorKVRecordTags(s.natgwKV, accountID, "nat-", input.Resources,
+	return utils.MirrorKVRecordTags(context.Background(), s.natgwKV, accountID, "nat-", input.Resources,
 		func(r *NatGatewayRecord) *map[string]string { return &r.Tags },
 		utils.MergeTagsMut(input))
 }
@@ -549,7 +561,7 @@ func (s *NatGatewayServiceImpl) RemoveRecordTags(input *ec2.DeleteTagsInput, acc
 	if input == nil {
 		return nil
 	}
-	return utils.MirrorKVRecordTags(s.natgwKV, accountID, "nat-", input.Resources,
+	return utils.MirrorKVRecordTags(context.Background(), s.natgwKV, accountID, "nat-", input.Resources,
 		func(r *NatGatewayRecord) *map[string]string { return &r.Tags },
 		utils.RemoveTagsMut(input))
 }

spinifex/handlers/ec2/natgw/service_impl_test.go

@@ -1,8 +1,8 @@
 package handlers_ec2_natgw
 
 import (
-	"context"
 	"testing"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
@@ -12,7 +12,7 @@ import (
 	handlers_ec2_vpc "github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc"
 	"github.com/mulgadc/spinifex/spinifex/testutil"
 	"github.com/mulgadc/spinifex/spinifex/utils"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -24,34 +24,74 @@ func setupTestService(t *testing.T) *NatGatewayServiceImpl {
 	return svc
 }
 
-func setupTestServiceJS(t *testing.T) (*NatGatewayServiceImpl, nats.JetStreamContext) {
+func setupTestServiceJS(t *testing.T) (*NatGatewayServiceImpl, jetstream.JetStream) {
 	t.Helper()
-	_, nc, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 
 	// Seed VPC KV
-	testutil.SeedKV(t, js, handlers_ec2_vpc.KVBucketVPCs, map[string][]byte{
+	seedKV(t, js, handlers_ec2_vpc.KVBucketVPCs, map[string][]byte{
 		utils.AccountKey(testAccountID, "vpc-test1"): []byte(`{"vpc_id":"vpc-test1","cidr_block":"10.0.0.0/16","state":"available"}`),
 	})
 
 	// Seed subnet KV (public subnet)
-	testutil.SeedKV(t, js, handlers_ec2_vpc.KVBucketSubnets, map[string][]byte{
+	seedKV(t, js, handlers_ec2_vpc.KVBucketSubnets, map[string][]byte{
 		utils.AccountKey(testAccountID, "subnet-pub1"): []byte(`{"subnet_id":"subnet-pub1","vpc_id":"vpc-test1","cidr_block":"10.0.1.0/24","state":"available","map_public_ip_on_launch":true}`),
 	})
 
 	// Seed EIP KV
-	testutil.SeedKV(t, js, handlers_ec2_eip.KVBucketEIPs, map[string][]byte{
+	seedKV(t, js, handlers_ec2_eip.KVBucketEIPs, map[string][]byte{
 		utils.AccountKey(testAccountID, "eipalloc-test1"): []byte(`{"allocation_id":"eipalloc-test1","public_ip":"203.0.113.50","state":"allocated"}`),
 		utils.AccountKey(testAccountID, "eipalloc-used"):  []byte(`{"allocation_id":"eipalloc-used","public_ip":"203.0.113.51","state":"associated","association_id":"eipassoc-xxx"}`),
 	})
 
-	svc, err := NewNatGatewayServiceImplWithNATS(nc)
+	svc, err := NewNatGatewayServiceImplWithNATS(t.Context(), nc)
 	require.NoError(t, err)
 	return svc, js
 }
 
+func seedKV(t *testing.T, js jetstream.JetStream, bucket string, entries map[string][]byte) jetstream.KeyValue {
+	t.Helper()
+	kv, err := js.CreateOrUpdateKeyValue(t.Context(), jetstream.KeyValueConfig{Bucket: bucket, History: 1})
+	require.NoError(t, err)
+	for key, value := range entries {
+		_, err := kv.Put(t.Context(), key, value)
+		require.NoError(t, err)
+	}
+	return kv
+}
+
+func TestDeletedBucketTTL(t *testing.T) {
+	_, js := setupTestServiceJS(t)
+	kv, err := js.KeyValue(t.Context(), KVBucketDeletedNatGateways)
+	require.NoError(t, err)
+	status, err := kv.Status(t.Context())
+	require.NoError(t, err)
+	assert.Equal(t, time.Hour, status.TTL())
+}
+
+func TestGetOrCreateDeletedBucketPreservesExistingConfiguration(t *testing.T) {
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	const existingTTL = 2 * time.Hour
+
+	_, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{
+		Bucket:  KVBucketDeletedNatGateways,
+		History: 1,
+		TTL:     existingTTL,
+	})
+	require.NoError(t, err)
+
+	kv, err := getOrCreateDeletedBucket(t.Context(), js)
+	require.NoError(t, err)
+	status, err := kv.Status(t.Context())
+	require.NoError(t, err)
+	require.Equal(t, existingTTL, status.TTL())
+}
+
 func TestCreateNatGateway(t *testing.T) {
 	svc := setupTestService(t)
-	out, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	out, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-test1"),
 	}, testAccountID)
@@ -72,7 +112,7 @@ func TestCreateNatGateway(t *testing.T) {
 // makes the tag-filtered describe return empty, leaking the NAT-GW EIP.
 func TestCreateNatGateway_PersistsTagsForTagFilterDiscovery(t *testing.T) {
 	svc := setupTestService(t)
-	_, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	_, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-test1"),
 		TagSpecifications: []*ec2.TagSpecification{{
@@ -85,7 +125,7 @@ func TestCreateNatGateway_PersistsTagsForTagFilterDiscovery(t *testing.T) {
 	}, testAccountID)
 	require.NoError(t, err)
 
-	out, err := svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err := svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		Filter: []*ec2.Filter{
 			{Name: aws.String("tag:spinifex:eks-cluster"), Values: aws.StringSlice([]string{"alpha"})},
 			{Name: aws.String("tag:spinifex:eks-role"), Values: aws.StringSlice([]string{"cp-natgw"})},
@@ -100,7 +140,7 @@ func TestCreateNatGateway_PersistsTagsForTagFilterDiscovery(t *testing.T) {
 
 func TestCreateNatGateway_SubnetNotFound(t *testing.T) {
 	svc := setupTestService(t)
-	_, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	_, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-nope"),
 		AllocationId: aws.String("eipalloc-test1"),
 	}, testAccountID)
@@ -109,7 +149,7 @@ func TestCreateNatGateway_SubnetNotFound(t *testing.T) {
 
 func TestCreateNatGateway_EIPNotFound(t *testing.T) {
 	svc := setupTestService(t)
-	_, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	_, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-nope"),
 	}, testAccountID)
@@ -118,7 +158,7 @@ func TestCreateNatGateway_EIPNotFound(t *testing.T) {
 
 func TestCreateNatGateway_EIPAlreadyAssociated(t *testing.T) {
 	svc := setupTestService(t)
-	_, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	_, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-used"),
 	}, testAccountID)
@@ -127,14 +167,14 @@ func TestCreateNatGateway_EIPAlreadyAssociated(t *testing.T) {
 
 func TestDeleteNatGateway(t *testing.T) {
 	svc := setupTestService(t)
-	createOut, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	createOut, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 	natgwID := *createOut.NatGateway.NatGatewayId
 
-	deleteOut, err := svc.DeleteNatGateway(context.Background(), &ec2.DeleteNatGatewayInput{
+	deleteOut, err := svc.DeleteNatGateway(t.Context(), &ec2.DeleteNatGatewayInput{
 		NatGatewayId: aws.String(natgwID),
 	}, testAccountID)
 	require.NoError(t, err)
@@ -147,13 +187,13 @@ func TestDeleteNatGateway(t *testing.T) {
 
 func TestDescribeNatGateways(t *testing.T) {
 	svc := setupTestService(t)
-	createOut, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	createOut, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	out, err := svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{}, testAccountID)
+	out, err := svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{}, testAccountID)
 	require.NoError(t, err)
 	assert.Len(t, out.NatGateways, 1)
 	assert.Equal(t, *createOut.NatGateway.NatGatewayId, *out.NatGateways[0].NatGatewayId)
@@ -161,21 +201,21 @@ func TestDescribeNatGateways(t *testing.T) {
 
 func TestDescribeNatGateways_ByID(t *testing.T) {
 	svc := setupTestService(t)
-	createOut, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	createOut, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 	natgwID := *createOut.NatGateway.NatGatewayId
 
-	out, err := svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err := svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		NatGatewayIds: []*string{aws.String(natgwID)},
 	}, testAccountID)
 	require.NoError(t, err)
 	assert.Len(t, out.NatGateways, 1)
 
 	// Non-existent ID
-	_, err = svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	_, err = svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		NatGatewayIds: []*string{aws.String("nat-nope")},
 	}, testAccountID)
 	assert.EqualError(t, err, awserrors.ErrorInvalidNatGatewayIDNotFound)
@@ -183,7 +223,7 @@ func TestDescribeNatGateways_ByID(t *testing.T) {
 
 func TestDescribeNatGateways_FilterByState(t *testing.T) {
 	svc := setupTestService(t)
-	_, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	_, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-test1"),
 	}, testAccountID)
@@ -191,14 +231,14 @@ func TestDescribeNatGateways_FilterByState(t *testing.T) {
 
 	name := "state"
 	val := "available"
-	out, err := svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err := svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		Filter: []*ec2.Filter{{Name: &name, Values: []*string{&val}}},
 	}, testAccountID)
 	require.NoError(t, err)
 	assert.Len(t, out.NatGateways, 1)
 
 	val2 := "deleted"
-	out, err = svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err = svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		Filter: []*ec2.Filter{{Name: &name, Values: []*string{&val2}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -207,7 +247,7 @@ func TestDescribeNatGateways_FilterByState(t *testing.T) {
 
 func TestGetNatGateway(t *testing.T) {
 	svc := setupTestService(t)
-	createOut, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	createOut, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-test1"),
 	}, testAccountID)
@@ -230,12 +270,12 @@ func TestGetNatGateway_NotFound(t *testing.T) {
 func TestCreateNatGateway_MissingParams(t *testing.T) {
 	svc := setupTestService(t)
 
-	_, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	_, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		AllocationId: aws.String("eipalloc-test1"),
 	}, testAccountID)
 	assert.EqualError(t, err, awserrors.ErrorMissingParameter)
 
-	_, err = svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	_, err = svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId: aws.String("subnet-pub1"),
 	}, testAccountID)
 	assert.EqualError(t, err, awserrors.ErrorMissingParameter)
@@ -243,20 +283,20 @@ func TestCreateNatGateway_MissingParams(t *testing.T) {
 
 func TestDescribeNatGateways_DeletedGatewayVisible(t *testing.T) {
 	svc := setupTestService(t)
-	createOut, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	createOut, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 	natgwID := *createOut.NatGateway.NatGatewayId
 
-	_, err = svc.DeleteNatGateway(context.Background(), &ec2.DeleteNatGatewayInput{
+	_, err = svc.DeleteNatGateway(t.Context(), &ec2.DeleteNatGatewayInput{
 		NatGatewayId: aws.String(natgwID),
 	}, testAccountID)
 	require.NoError(t, err)
 
 	// Describe by ID should return the deleted gateway with state=deleted
-	out, err := svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err := svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		NatGatewayIds: []*string{aws.String(natgwID)},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -264,12 +304,12 @@ func TestDescribeNatGateways_DeletedGatewayVisible(t *testing.T) {
 	assert.Equal(t, "deleted", *out.NatGateways[0].State)
 
 	// Describe without filter should NOT include deleted gateways
-	out, err = svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{}, testAccountID)
+	out, err = svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{}, testAccountID)
 	require.NoError(t, err)
 	assert.Empty(t, out.NatGateways)
 
 	// Non-existent ID should still error
-	_, err = svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	_, err = svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		NatGatewayIds: []*string{aws.String("nat-never-existed")},
 	}, testAccountID)
 	assert.EqualError(t, err, awserrors.ErrorInvalidNatGatewayIDNotFound)
@@ -277,7 +317,7 @@ func TestDescribeNatGateways_DeletedGatewayVisible(t *testing.T) {
 
 func TestDescribeNatGateways_FilterByNatGatewayId(t *testing.T) {
 	svc := setupTestService(t)
-	createOut, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	createOut, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-test1"),
 	}, testAccountID)
@@ -285,7 +325,7 @@ func TestDescribeNatGateways_FilterByNatGatewayId(t *testing.T) {
 	natgwID := *createOut.NatGateway.NatGatewayId
 
 	// Exact match
-	out, err := svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err := svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		Filter: []*ec2.Filter{{Name: aws.String("nat-gateway-id"), Values: []*string{aws.String(natgwID)}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -293,14 +333,14 @@ func TestDescribeNatGateways_FilterByNatGatewayId(t *testing.T) {
 	assert.Equal(t, natgwID, *out.NatGateways[0].NatGatewayId)
 
 	// Non-match
-	out, err = svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err = svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		Filter: []*ec2.Filter{{Name: aws.String("nat-gateway-id"), Values: []*string{aws.String("nat-000000")}}},
 	}, testAccountID)
 	require.NoError(t, err)
 	assert.Empty(t, out.NatGateways)
 
 	// Wildcard
-	out, err = svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err = svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		Filter: []*ec2.Filter{{Name: aws.String("nat-gateway-id"), Values: []*string{aws.String("nat-*")}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -309,21 +349,21 @@ func TestDescribeNatGateways_FilterByNatGatewayId(t *testing.T) {
 
 func TestDescribeNatGateways_FilterBySubnetId(t *testing.T) {
 	svc := setupTestService(t)
-	_, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	_, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
 	// Exact match
-	out, err := svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err := svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		Filter: []*ec2.Filter{{Name: aws.String("subnet-id"), Values: []*string{aws.String("subnet-pub1")}}},
 	}, testAccountID)
 	require.NoError(t, err)
 	assert.Len(t, out.NatGateways, 1)
 
 	// Non-match
-	out, err = svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err = svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		Filter: []*ec2.Filter{{Name: aws.String("subnet-id"), Values: []*string{aws.String("subnet-000000")}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -332,7 +372,7 @@ func TestDescribeNatGateways_FilterBySubnetId(t *testing.T) {
 
 func TestDescribeNatGateways_FilterByVpcId(t *testing.T) {
 	svc := setupTestService(t)
-	_, err := svc.CreateNatGateway(context.Background(), &ec2.CreateNatGatewayInput{
+	_, err := svc.CreateNatGateway(t.Context(), &ec2.CreateNatGatewayInput{
 		SubnetId:     aws.String("subnet-pub1"),
 		AllocationId: aws.String("eipalloc-test1"),
 	}, testAccountID)
@@ -340,14 +380,14 @@ func TestDescribeNatGateways_FilterByVpcId(t *testing.T) {
 
 	name := "vpc-id"
 	val := "vpc-test1"
-	out, err := svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err := svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		Filter: []*ec2.Filter{{Name: &name, Values: []*string{&val}}},
 	}, testAccountID)
 	require.NoError(t, err)
 	assert.Len(t, out.NatGateways, 1)
 
 	val2 := "vpc-nope"
-	out, err = svc.DescribeNatGateways(context.Background(), &ec2.DescribeNatGatewaysInput{
+	out, err = svc.DescribeNatGateways(t.Context(), &ec2.DescribeNatGatewaysInput{
 		Filter: []*ec2.Filter{{Name: &name, Values: []*string{&val2}}},
 	}, testAccountID)
 	require.NoError(t, err)

spinifex/handlers/ec2/placementgroup/service_impl.go

@@ -14,9 +14,11 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Ensure PlacementGroupServiceImpl implements PlacementGroupService.
@@ -46,21 +48,21 @@ type PlacementGroupRecord struct {
 type PlacementGroupServiceImpl struct {
 	config   *config.Config
 	natsConn *nats.Conn
-	kv       nats.KeyValue
+	kv       jetstream.KeyValue
 }
 
 // NewPlacementGroupServiceImplWithNATS creates a placement group service with NATS JetStream.
-func NewPlacementGroupServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*PlacementGroupServiceImpl, error) {
-	js, err := natsConn.JetStream()
+func NewPlacementGroupServiceImplWithNATS(ctx context.Context, cfg *config.Config, natsConn *nats.Conn) (*PlacementGroupServiceImpl, error) {
+	js, err := jetstream.New(natsConn)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get JetStream context: %w", err)
 	}
 
-	kv, err := utils.GetOrCreateKVBucket(js, KVBucketPlacementGroups, 10)
+	kv, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketPlacementGroups, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketPlacementGroups, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketPlacementGroups, kv, KVBucketPlacementGroupsVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketPlacementGroups, kv, KVBucketPlacementGroupsVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketPlacementGroups, err)
 	}
 
@@ -113,7 +115,7 @@ func (s *PlacementGroupServiceImpl) CreatePlacementGroup(ctx context.Context, in
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 	// Atomic create-if-not-exists to prevent TOCTOU race on duplicate names
-	if _, err := s.kv.Create(key, data); err != nil {
+	if _, err := s.kv.Create(ctx, key, data); err != nil {
 		// Create fails if key already exists
 		return nil, errors.New(awserrors.ErrorInvalidPlacementGroupDuplicate)
 	}
@@ -134,7 +136,7 @@ func (s *PlacementGroupServiceImpl) DeletePlacementGroup(ctx context.Context, in
 	groupName := *input.GroupName
 	key := utils.AccountKey(accountID, groupName)
 
-	entry, err := s.kv.Get(key)
+	entry, err := s.kv.Get(ctx, key)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidPlacementGroupUnknown)
 	}
@@ -153,7 +155,7 @@ func (s *PlacementGroupServiceImpl) DeletePlacementGroup(ctx context.Context, in
 		return nil, errors.New(awserrors.ErrorInvalidPlacementGroupInUse)
 	}
 
-	if err := s.kv.Delete(key); err != nil {
+	if err := s.kv.Delete(ctx, key); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -195,8 +197,8 @@ func (s *PlacementGroupServiceImpl) DescribePlacementGroups(ctx context.Context,
 	}
 
 	prefix := accountID + "."
-	keys, err := s.kv.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	keys, err := s.kv.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -209,7 +211,7 @@ func (s *PlacementGroupServiceImpl) DescribePlacementGroups(ctx context.Context,
 			continue
 		}
 
-		entry, err := s.kv.Get(k)
+		entry, err := s.kv.Get(ctx, k)
 		if err != nil {
 			slog.WarnContext(ctx, "Failed to get placement group record", "key", k, "error", err)
 			continue
@@ -313,9 +315,9 @@ func pgMatchesAnyTag(tags map[string]string, values []string, field func(k, v st
 
 // GetPlacementGroupRecord reads a placement group record from KV with its revision for CAS operations.
 // Returns the record and the KV entry (for revision). Exported for use by gateway spread routing.
-func (s *PlacementGroupServiceImpl) GetPlacementGroupRecord(accountID, groupName string) (*PlacementGroupRecord, nats.KeyValueEntry, error) {
+func (s *PlacementGroupServiceImpl) GetPlacementGroupRecord(ctx context.Context, accountID, groupName string) (*PlacementGroupRecord, jetstream.KeyValueEntry, error) {
 	key := utils.AccountKey(accountID, groupName)
-	entry, err := s.kv.Get(key)
+	entry, err := s.kv.Get(ctx, key)
 	if err != nil {
 		return nil, nil, errors.New(awserrors.ErrorInvalidPlacementGroupUnknown)
 	}
@@ -330,13 +332,13 @@ func (s *PlacementGroupServiceImpl) GetPlacementGroupRecord(accountID, groupName
 
 // UpdatePlacementGroupRecord writes a placement group record using CAS (optimistic concurrency).
 // Returns nil on success or the error on CAS conflict.
-func (s *PlacementGroupServiceImpl) UpdatePlacementGroupRecord(accountID, groupName string, record *PlacementGroupRecord, revision uint64) error {
+func (s *PlacementGroupServiceImpl) UpdatePlacementGroupRecord(ctx context.Context, accountID, groupName string, record *PlacementGroupRecord, revision uint64) error {
 	key := utils.AccountKey(accountID, groupName)
 	data, err := json.Marshal(record)
 	if err != nil {
 		return errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.kv.Update(key, data, revision); err != nil {
+	if _, err := s.kv.Update(ctx, key, data, revision); err != nil {
 		return err
 	}
 	return nil
@@ -352,7 +354,7 @@ func (s *PlacementGroupServiceImpl) ReserveSpreadNodes(ctx context.Context, inpu
 	}
 
 	for attempt := range maxCASRetries {
-		record, entry, err := s.GetPlacementGroupRecord(accountID, input.GroupName)
+		record, entry, err := s.GetPlacementGroupRecord(ctx, accountID, input.GroupName)
 		if err != nil {
 			return nil, err
 		}
@@ -392,7 +394,7 @@ func (s *PlacementGroupServiceImpl) ReserveSpreadNodes(ctx context.Context, inpu
 		}
 
 		// CAS write — retry on conflict
-		if err := s.UpdatePlacementGroupRecord(accountID, input.GroupName, record, entry.Revision()); err != nil {
+		if err := s.UpdatePlacementGroupRecord(ctx, accountID, input.GroupName, record, entry.Revision()); err != nil {
 			slog.DebugContext(ctx, "ReserveSpreadNodes: CAS conflict, retrying", "attempt", attempt, "err", err)
 			continue
 		}
@@ -413,7 +415,7 @@ func (s *PlacementGroupServiceImpl) FinalizeSpreadInstances(ctx context.Context,
 	}
 
 	for attempt := range maxCASRetries {
-		record, entry, err := s.GetPlacementGroupRecord(accountID, input.GroupName)
+		record, entry, err := s.GetPlacementGroupRecord(ctx, accountID, input.GroupName)
 		if err != nil {
 			return nil, err
 		}
@@ -421,7 +423,7 @@ func (s *PlacementGroupServiceImpl) FinalizeSpreadInstances(ctx context.Context,
 		// Replace placeholder entries with actual instance IDs per node
 		maps.Copy(record.NodeInstances, input.NodeInstances)
 
-		if err := s.UpdatePlacementGroupRecord(accountID, input.GroupName, record, entry.Revision()); err != nil {
+		if err := s.UpdatePlacementGroupRecord(ctx, accountID, input.GroupName, record, entry.Revision()); err != nil {
 			slog.DebugContext(ctx, "FinalizeSpreadInstances: CAS conflict, retrying", "attempt", attempt, "err", err)
 			continue
 		}
@@ -447,7 +449,7 @@ func (s *PlacementGroupServiceImpl) ReleaseSpreadNodes(ctx context.Context, inpu
 	}
 
 	for attempt := range maxCASRetries {
-		record, entry, err := s.GetPlacementGroupRecord(accountID, input.GroupName)
+		record, entry, err := s.GetPlacementGroupRecord(ctx, accountID, input.GroupName)
 		if err != nil {
 			return nil, err
 		}
@@ -460,7 +462,7 @@ func (s *PlacementGroupServiceImpl) ReleaseSpreadNodes(ctx context.Context, inpu
 			}
 		}
 
-		if err := s.UpdatePlacementGroupRecord(accountID, input.GroupName, record, entry.Revision()); err != nil {
+		if err := s.UpdatePlacementGroupRecord(ctx, accountID, input.GroupName, record, entry.Revision()); err != nil {
 			slog.DebugContext(ctx, "ReleaseSpreadNodes: CAS conflict, retrying", "attempt", attempt, "err", err)
 			continue
 		}
@@ -482,7 +484,7 @@ func (s *PlacementGroupServiceImpl) RemoveInstance(ctx context.Context, input *R
 	}
 
 	for attempt := range maxCASRetries {
-		record, entry, lookupErr := s.GetPlacementGroupRecord(accountID, input.GroupName)
+		record, entry, lookupErr := s.GetPlacementGroupRecord(ctx, accountID, input.GroupName)
 		if lookupErr != nil {
 			// Group may have been deleted already — treat as success
 			slog.DebugContext(ctx, "RemoveInstance: group not found, treating as success", "groupName", input.GroupName)
@@ -509,7 +511,7 @@ func (s *PlacementGroupServiceImpl) RemoveInstance(ctx context.Context, input *R
 			record.NodeInstances[input.NodeName] = filtered
 		}
 
-		if err := s.UpdatePlacementGroupRecord(accountID, input.GroupName, record, entry.Revision()); err != nil {
+		if err := s.UpdatePlacementGroupRecord(ctx, accountID, input.GroupName, record, entry.Revision()); err != nil {
 			slog.DebugContext(ctx, "RemoveInstance: CAS conflict, retrying", "attempt", attempt, "err", err)
 			continue
 		}
@@ -531,7 +533,7 @@ func (s *PlacementGroupServiceImpl) ReserveClusterNode(ctx context.Context, inpu
 	}
 
 	for attempt := range maxCASRetries {
-		record, entry, err := s.GetPlacementGroupRecord(accountID, input.GroupName)
+		record, entry, err := s.GetPlacementGroupRecord(ctx, accountID, input.GroupName)
 		if err != nil {
 			return nil, err
 		}
@@ -555,7 +557,7 @@ func (s *PlacementGroupServiceImpl) ReserveClusterNode(ctx context.Context, inpu
 		targetNode := input.EligibleNodes[0] // first = highest capacity (sorted desc by caller)
 		record.NodeInstances[targetNode] = []string{}
 
-		if err := s.UpdatePlacementGroupRecord(accountID, input.GroupName, record, entry.Revision()); err != nil {
+		if err := s.UpdatePlacementGroupRecord(ctx, accountID, input.GroupName, record, entry.Revision()); err != nil {
 			slog.DebugContext(ctx, "ReserveClusterNode: CAS conflict, retrying", "attempt", attempt, "err", err)
 			continue
 		}
@@ -576,7 +578,7 @@ func (s *PlacementGroupServiceImpl) FinalizeClusterInstances(ctx context.Context
 	}
 
 	for attempt := range maxCASRetries {
-		record, entry, err := s.GetPlacementGroupRecord(accountID, input.GroupName)
+		record, entry, err := s.GetPlacementGroupRecord(ctx, accountID, input.GroupName)
 		if err != nil {
 			return nil, err
 		}
@@ -586,7 +588,7 @@ func (s *PlacementGroupServiceImpl) FinalizeClusterInstances(ctx context.Context
 			record.NodeInstances[node] = append(record.NodeInstances[node], ids...)
 		}
 
-		if err := s.UpdatePlacementGroupRecord(accountID, input.GroupName, record, entry.Revision()); err != nil {
+		if err := s.UpdatePlacementGroupRecord(ctx, accountID, input.GroupName, record, entry.Revision()); err != nil {
 			slog.DebugContext(ctx, "FinalizeClusterInstances: CAS conflict, retrying", "attempt", attempt, "err", err)
 			continue
 		}

spinifex/handlers/ec2/placementgroup/service_impl_test.go

@@ -18,7 +18,7 @@ func setupTestService(t *testing.T) *PlacementGroupServiceImpl {
 	t.Helper()
 	_, nc, _ := testutil.StartTestJetStream(t)
 
-	svc, err := NewPlacementGroupServiceImplWithNATS(nil, nc)
+	svc, err := NewPlacementGroupServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 	return svc
 }
@@ -162,10 +162,10 @@ func TestDeletePlacementGroup_InUse(t *testing.T) {
 	createTestGroup(t, svc, "in-use-group", "spread")
 
 	// Simulate instances by updating the record directly
-	record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "in-use-group")
+	record, entry, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "in-use-group")
 	require.NoError(t, err)
 	record.NodeInstances["node1"] = []string{"i-123"}
-	err = svc.UpdatePlacementGroupRecord(testAccountID, "in-use-group", record, entry.Revision())
+	err = svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "in-use-group", record, entry.Revision())
 	require.NoError(t, err)
 
 	_, err = svc.DeletePlacementGroup(context.Background(), &ec2.DeletePlacementGroupInput{
@@ -315,18 +315,18 @@ func TestGetAndUpdatePlacementGroupRecord(t *testing.T) {
 	createTestGroup(t, svc, "cas-group", "spread")
 
 	// Get with revision
-	record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "cas-group")
+	record, entry, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "cas-group")
 	require.NoError(t, err)
 	assert.Equal(t, "spread", record.Strategy)
 	assert.Empty(t, record.NodeInstances)
 
 	// Update with CAS
 	record.NodeInstances["node1"] = []string{"i-abc"}
-	err = svc.UpdatePlacementGroupRecord(testAccountID, "cas-group", record, entry.Revision())
+	err = svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "cas-group", record, entry.Revision())
 	require.NoError(t, err)
 
 	// Verify update
-	record2, _, err := svc.GetPlacementGroupRecord(testAccountID, "cas-group")
+	record2, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "cas-group")
 	require.NoError(t, err)
 	assert.Equal(t, []string{"i-abc"}, record2.NodeInstances["node1"])
 }
@@ -336,25 +336,25 @@ func TestUpdatePlacementGroupRecord_CASConflict(t *testing.T) {
 	createTestGroup(t, svc, "conflict-group", "spread")
 
 	// Get the record twice (same revision)
-	record1, entry1, err := svc.GetPlacementGroupRecord(testAccountID, "conflict-group")
+	record1, entry1, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "conflict-group")
 	require.NoError(t, err)
-	record2, _, err := svc.GetPlacementGroupRecord(testAccountID, "conflict-group")
+	record2, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "conflict-group")
 	require.NoError(t, err)
 
 	// First update succeeds
 	record1.NodeInstances["node1"] = []string{"i-111"}
-	err = svc.UpdatePlacementGroupRecord(testAccountID, "conflict-group", record1, entry1.Revision())
+	err = svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "conflict-group", record1, entry1.Revision())
 	require.NoError(t, err)
 
 	// Second update with stale revision fails
 	record2.NodeInstances["node2"] = []string{"i-222"}
-	err = svc.UpdatePlacementGroupRecord(testAccountID, "conflict-group", record2, entry1.Revision())
+	err = svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "conflict-group", record2, entry1.Revision())
 	require.Error(t, err)
 }
 
 func TestGetPlacementGroupRecord_NotFound(t *testing.T) {
 	svc := setupTestService(t)
-	_, _, err := svc.GetPlacementGroupRecord(testAccountID, "nonexistent")
+	_, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "nonexistent")
 	require.Error(t, err)
 	assert.Equal(t, awserrors.ErrorInvalidPlacementGroupUnknown, err.Error())
 }
@@ -375,7 +375,7 @@ func TestReserveSpreadNodes_Success(t *testing.T) {
 	assert.Len(t, out.ReservedNodes, 3)
 
 	// Verify placeholders in record
-	record, _, err := svc.GetPlacementGroupRecord(testAccountID, "reserve-group")
+	record, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "reserve-group")
 	require.NoError(t, err)
 	assert.Len(t, record.NodeInstances, 3)
 	for _, node := range out.ReservedNodes {
@@ -389,10 +389,10 @@ func TestReserveSpreadNodes_ExcludesOccupiedNodes(t *testing.T) {
 	createTestGroup(t, svc, "occupied-group", "spread")
 
 	// Pre-occupy node-a
-	record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "occupied-group")
+	record, entry, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "occupied-group")
 	require.NoError(t, err)
 	record.NodeInstances["node-a"] = []string{"i-existing"}
-	require.NoError(t, svc.UpdatePlacementGroupRecord(testAccountID, "occupied-group", record, entry.Revision()))
+	require.NoError(t, svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "occupied-group", record, entry.Revision()))
 
 	out, err := svc.ReserveSpreadNodes(context.Background(), &ReserveSpreadNodesInput{
 		GroupName:     "occupied-group",
@@ -413,11 +413,11 @@ func TestReserveSpreadNodes_InsufficientNodes(t *testing.T) {
 	createTestGroup(t, svc, "insufficient-group", "spread")
 
 	// Pre-occupy node-a and node-b
-	record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "insufficient-group")
+	record, entry, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "insufficient-group")
 	require.NoError(t, err)
 	record.NodeInstances["node-a"] = []string{"i-1"}
 	record.NodeInstances["node-b"] = []string{"i-2"}
-	require.NoError(t, svc.UpdatePlacementGroupRecord(testAccountID, "insufficient-group", record, entry.Revision()))
+	require.NoError(t, svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "insufficient-group", record, entry.Revision()))
 
 	_, err = svc.ReserveSpreadNodes(context.Background(), &ReserveSpreadNodesInput{
 		GroupName:     "insufficient-group",
@@ -481,7 +481,7 @@ func TestFinalizeSpreadInstances_Success(t *testing.T) {
 	require.NoError(t, err)
 
 	// Verify
-	record, _, err := svc.GetPlacementGroupRecord(testAccountID, "finalize-group")
+	record, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "finalize-group")
 	require.NoError(t, err)
 	assert.Equal(t, []string{"i-aaa"}, record.NodeInstances["node-a"])
 	assert.Equal(t, []string{"i-bbb"}, record.NodeInstances["node-b"])
@@ -510,7 +510,7 @@ func TestReleaseSpreadNodes_Success(t *testing.T) {
 	require.NoError(t, err)
 
 	// Verify only node-a remains
-	record, _, err := svc.GetPlacementGroupRecord(testAccountID, "release-group")
+	record, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "release-group")
 	require.NoError(t, err)
 	assert.Len(t, record.NodeInstances, 1)
 	_, ok := record.NodeInstances["node-a"]
@@ -538,7 +538,7 @@ func TestReleaseSpreadNodes_AllNodes(t *testing.T) {
 	require.NoError(t, err)
 
 	// Verify empty
-	record, _, err := svc.GetPlacementGroupRecord(testAccountID, "release-all-group")
+	record, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "release-all-group")
 	require.NoError(t, err)
 	assert.Empty(t, record.NodeInstances)
 }
@@ -550,11 +550,11 @@ func TestRemoveInstance_Success(t *testing.T) {
 	createTestGroup(t, svc, "remove-inst-group", "spread")
 
 	// Add instances to two nodes
-	record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "remove-inst-group")
+	record, entry, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "remove-inst-group")
 	require.NoError(t, err)
 	record.NodeInstances["node-a"] = []string{"i-aaa"}
 	record.NodeInstances["node-b"] = []string{"i-bbb"}
-	require.NoError(t, svc.UpdatePlacementGroupRecord(testAccountID, "remove-inst-group", record, entry.Revision()))
+	require.NoError(t, svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "remove-inst-group", record, entry.Revision()))
 
 	// Remove i-aaa from node-a
 	_, err = svc.RemoveInstance(context.Background(), &RemoveInstanceInput{
@@ -565,7 +565,7 @@ func TestRemoveInstance_Success(t *testing.T) {
 	require.NoError(t, err)
 
 	// Verify node-a is removed (was the only instance), node-b remains
-	record, _, err = svc.GetPlacementGroupRecord(testAccountID, "remove-inst-group")
+	record, _, err = svc.GetPlacementGroupRecord(t.Context(), testAccountID, "remove-inst-group")
 	require.NoError(t, err)
 	assert.Len(t, record.NodeInstances, 1)
 	_, hasA := record.NodeInstances["node-a"]
@@ -605,10 +605,10 @@ func TestRemoveInstance_MultipleInstancesOnNode(t *testing.T) {
 	createTestGroup(t, svc, "multi-inst-group", "cluster")
 
 	// Add multiple instances to same node (cluster scenario)
-	record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "multi-inst-group")
+	record, entry, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "multi-inst-group")
 	require.NoError(t, err)
 	record.NodeInstances["node-a"] = []string{"i-111", "i-222", "i-333"}
-	require.NoError(t, svc.UpdatePlacementGroupRecord(testAccountID, "multi-inst-group", record, entry.Revision()))
+	require.NoError(t, svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "multi-inst-group", record, entry.Revision()))
 
 	// Remove i-222
 	_, err = svc.RemoveInstance(context.Background(), &RemoveInstanceInput{
@@ -619,7 +619,7 @@ func TestRemoveInstance_MultipleInstancesOnNode(t *testing.T) {
 	require.NoError(t, err)
 
 	// Verify i-111 and i-333 remain
-	record, _, err = svc.GetPlacementGroupRecord(testAccountID, "multi-inst-group")
+	record, _, err = svc.GetPlacementGroupRecord(t.Context(), testAccountID, "multi-inst-group")
 	require.NoError(t, err)
 	assert.Equal(t, []string{"i-111", "i-333"}, record.NodeInstances["node-a"])
 }
@@ -639,7 +639,7 @@ func TestReserveClusterNode_EmptyGroup(t *testing.T) {
 	assert.Equal(t, "node-a", out.TargetNode)
 
 	// Verify placeholder written
-	record, _, err := svc.GetPlacementGroupRecord(testAccountID, "cluster-empty")
+	record, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "cluster-empty")
 	require.NoError(t, err)
 	assert.Len(t, record.NodeInstances, 1)
 	_, ok := record.NodeInstances["node-a"]
@@ -651,10 +651,10 @@ func TestReserveClusterNode_ExistingNode(t *testing.T) {
 	createTestGroup(t, svc, "cluster-existing", "cluster")
 
 	// Pre-populate with instances on node-b
-	record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "cluster-existing")
+	record, entry, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "cluster-existing")
 	require.NoError(t, err)
 	record.NodeInstances["node-b"] = []string{"i-existing"}
-	require.NoError(t, svc.UpdatePlacementGroupRecord(testAccountID, "cluster-existing", record, entry.Revision()))
+	require.NoError(t, svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "cluster-existing", record, entry.Revision()))
 
 	// Even though node-a is in eligible list, should return existing node-b
 	out, err := svc.ReserveClusterNode(context.Background(), &ReserveClusterNodeInput{
@@ -731,7 +731,7 @@ func TestFinalizeClusterInstances_Success(t *testing.T) {
 	require.NoError(t, err)
 
 	// Verify instances recorded
-	record, _, err := svc.GetPlacementGroupRecord(testAccountID, "cluster-finalize")
+	record, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "cluster-finalize")
 	require.NoError(t, err)
 	assert.Equal(t, []string{"i-c1", "i-c2", "i-c3"}, record.NodeInstances["node-a"])
 }
@@ -741,10 +741,10 @@ func TestFinalizeClusterInstances_AppendsToExisting(t *testing.T) {
 	createTestGroup(t, svc, "cluster-append", "cluster")
 
 	// Pre-populate with existing instances
-	record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "cluster-append")
+	record, entry, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "cluster-append")
 	require.NoError(t, err)
 	record.NodeInstances["node-a"] = []string{"i-existing1", "i-existing2"}
-	require.NoError(t, svc.UpdatePlacementGroupRecord(testAccountID, "cluster-append", record, entry.Revision()))
+	require.NoError(t, svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "cluster-append", record, entry.Revision()))
 
 	// Finalize with new instances (should append, not overwrite)
 	_, err = svc.FinalizeClusterInstances(context.Background(), &FinalizeClusterInstancesInput{
@@ -756,7 +756,7 @@ func TestFinalizeClusterInstances_AppendsToExisting(t *testing.T) {
 	require.NoError(t, err)
 
 	// Verify all instances present
-	record, _, err = svc.GetPlacementGroupRecord(testAccountID, "cluster-append")
+	record, _, err = svc.GetPlacementGroupRecord(t.Context(), testAccountID, "cluster-append")
 	require.NoError(t, err)
 	assert.Equal(t, []string{"i-existing1", "i-existing2", "i-new1", "i-new2"}, record.NodeInstances["node-a"])
 }

spinifex/handlers/ec2/routetable/lifecycle_invariants_test.go

@@ -1,7 +1,6 @@
 package handlers_ec2_routetable
 
 import (
-	"context"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -19,7 +18,7 @@ import (
 func TestRLC1_RouteTableDeleteNotFoundOnAbsent(t *testing.T) {
 	svc := setupTestService(t)
 
-	_, err := svc.DeleteRouteTable(context.Background(), &ec2.DeleteRouteTableInput{
+	_, err := svc.DeleteRouteTable(t.Context(), &ec2.DeleteRouteTableInput{
 		RouteTableId: aws.String("rtb-absent00000000"),
 	}, testAccountID)
 

spinifex/handlers/ec2/routetable/record_tags_test.go

@@ -1,7 +1,6 @@
 package handlers_ec2_routetable
 
 import (
-	"context"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -31,7 +30,7 @@ func TestRouteTableRecordTagsMirror(t *testing.T) {
 		},
 	}, testAccountID))
 
-	out, err := svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err := svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		RouteTableIds: []*string{aws.String(rtbID)},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -48,7 +47,7 @@ func TestRouteTableRecordTagsMirror(t *testing.T) {
 		},
 	}, testAccountID))
 
-	out, err = svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err = svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		RouteTableIds: []*string{aws.String(rtbID)},
 	}, testAccountID)
 	require.NoError(t, err)

spinifex/handlers/ec2/routetable/service_impl.go

@@ -17,9 +17,11 @@ import (
 	handlers_ec2_igw "github.com/mulgadc/spinifex/spinifex/handlers/ec2/igw"
 	handlers_ec2_natgw "github.com/mulgadc/spinifex/spinifex/handlers/ec2/natgw"
 	handlers_ec2_vpc "github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Ensure RouteTableServiceImpl implements RouteTableService.
@@ -28,45 +30,45 @@ var _ RouteTableService = (*RouteTableServiceImpl)(nil)
 // RouteTableServiceImpl implements Route Table operations with NATS JetStream persistence.
 type RouteTableServiceImpl struct {
 	config   *config.Config
-	rtbKV    nats.KeyValue
-	vpcKV    nats.KeyValue
-	igwKV    nats.KeyValue
-	subnetKV nats.KeyValue
-	natgwKV  nats.KeyValue
+	rtbKV    jetstream.KeyValue
+	vpcKV    jetstream.KeyValue
+	igwKV    jetstream.KeyValue
+	subnetKV jetstream.KeyValue
+	natgwKV  jetstream.KeyValue
 	natsConn *nats.Conn
 }
 
 // NewRouteTableServiceImplWithNATS creates a Route Table service with NATS JetStream for persistence.
-func NewRouteTableServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*RouteTableServiceImpl, error) {
-	js, err := natsConn.JetStream()
+func NewRouteTableServiceImplWithNATS(ctx context.Context, cfg *config.Config, natsConn *nats.Conn) (*RouteTableServiceImpl, error) {
+	js, err := jetstream.New(natsConn)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get JetStream context: %w", err)
 	}
 
-	rtbKV, err := utils.GetOrCreateKVBucket(js, KVBucketRouteTables, 10)
+	rtbKV, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketRouteTables, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketRouteTables, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketRouteTables, rtbKV, KVBucketRouteTablesVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketRouteTables, rtbKV, KVBucketRouteTablesVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketRouteTables, err)
 	}
 
-	vpcKV, err := utils.GetOrCreateKVBucket(js, handlers_ec2_vpc.KVBucketVPCs, 10)
+	vpcKV, err := kvutil.GetOrCreateBucket(ctx, js, handlers_ec2_vpc.KVBucketVPCs, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get VPC KV bucket: %w", err)
 	}
 
-	igwKV, err := utils.GetOrCreateKVBucket(js, handlers_ec2_igw.KVBucketIGW, 10)
+	igwKV, err := kvutil.GetOrCreateBucket(ctx, js, handlers_ec2_igw.KVBucketIGW, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get IGW KV bucket: %w", err)
 	}
 
-	subnetKV, err := utils.GetOrCreateKVBucket(js, handlers_ec2_vpc.KVBucketSubnets, 10)
+	subnetKV, err := kvutil.GetOrCreateBucket(ctx, js, handlers_ec2_vpc.KVBucketSubnets, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get subnet KV bucket: %w", err)
 	}
 
-	natgwKV, err := utils.GetOrCreateKVBucket(js, handlers_ec2_natgw.KVBucketNatGateways, 10)
+	natgwKV, err := kvutil.GetOrCreateBucket(ctx, js, handlers_ec2_natgw.KVBucketNatGateways, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get NAT Gateway KV bucket: %w", err)
 	}
@@ -86,9 +88,9 @@ func NewRouteTableServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (
 
 // getRouteTable retrieves a route table record from KV.
 func (s *RouteTableServiceImpl) getRouteTable(ctx context.Context, accountID, rtbID string) (*RouteTableRecord, error) {
-	entry, err := s.rtbKV.Get(utils.AccountKey(accountID, rtbID))
+	entry, err := s.rtbKV.Get(ctx, utils.AccountKey(accountID, rtbID))
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return nil, errors.New(awserrors.ErrorInvalidRouteTableIDNotFound)
 		}
 		slog.ErrorContext(ctx, "Failed to read route table from KV", "routeTableId", rtbID, "err", err)
@@ -109,7 +111,7 @@ func (s *RouteTableServiceImpl) putRouteTable(ctx context.Context, accountID str
 		slog.ErrorContext(ctx, "Failed to marshal route table record", "routeTableId", record.RouteTableId, "err", err)
 		return errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.rtbKV.Put(utils.AccountKey(accountID, record.RouteTableId), data); err != nil {
+	if _, err := s.rtbKV.Put(ctx, utils.AccountKey(accountID, record.RouteTableId), data); err != nil {
 		slog.ErrorContext(ctx, "Failed to write route table to KV", "routeTableId", record.RouteTableId, "err", err)
 		return errors.New(awserrors.ErrorServerInternal)
 	}
@@ -130,9 +132,9 @@ const rtbCASMaxRetries = 16
 func (s *RouteTableServiceImpl) mutateRouteTableCAS(ctx context.Context, accountID, rtbID string, mutate func(*RouteTableRecord) (bool, error)) error {
 	key := utils.AccountKey(accountID, rtbID)
 	for range rtbCASMaxRetries {
-		entry, err := s.rtbKV.Get(key)
+		entry, err := s.rtbKV.Get(ctx, key)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				return errors.New(awserrors.ErrorInvalidRouteTableIDNotFound)
 			}
 			slog.ErrorContext(ctx, "Failed to read route table from KV", "routeTableId", rtbID, "err", err)
@@ -158,8 +160,8 @@ func (s *RouteTableServiceImpl) mutateRouteTableCAS(ctx context.Context, account
 			slog.ErrorContext(ctx, "Failed to marshal route table record", "routeTableId", rtbID, "err", err)
 			return errors.New(awserrors.ErrorServerInternal)
 		}
-		if _, err := s.rtbKV.Update(key, data, entry.Revision()); err != nil {
-			if errors.Is(err, nats.ErrKeyExists) {
+		if _, err := s.rtbKV.Update(ctx, key, data, entry.Revision()); err != nil {
+			if errors.Is(err, jetstream.ErrKeyExists) {
 				continue // CAS conflict — another writer won, re-read and retry.
 			}
 			slog.ErrorContext(ctx, "Failed to write route table to KV", "routeTableId", rtbID, "err", err)
@@ -172,7 +174,7 @@ func (s *RouteTableServiceImpl) mutateRouteTableCAS(ctx context.Context, account
 
 // getVPCCidr looks up a VPC's CIDR block from the VPC KV bucket.
 func (s *RouteTableServiceImpl) getVPCCidr(ctx context.Context, accountID, vpcID string) (string, error) {
-	entry, err := s.vpcKV.Get(utils.AccountKey(accountID, vpcID))
+	entry, err := s.vpcKV.Get(ctx, utils.AccountKey(accountID, vpcID))
 	if err != nil {
 		return "", errors.New(awserrors.ErrorInvalidVpcIDNotFound)
 	}
@@ -219,9 +221,9 @@ func (s *RouteTableServiceImpl) subnetsImplicitlyOnMainRT(ctx context.Context, a
 	}
 
 	prefix := accountID + "."
-	keys, err := s.subnetKV.Keys()
+	keys, err := s.subnetKV.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, nil
 		}
 		return nil, fmt.Errorf("list subnet keys: %w", err)
@@ -231,9 +233,9 @@ func (s *RouteTableServiceImpl) subnetsImplicitlyOnMainRT(ctx context.Context, a
 		if key == utils.VersionKey || !strings.HasPrefix(key, prefix) {
 			continue
 		}
-		entry, err := s.subnetKV.Get(key)
+		entry, err := s.subnetKV.Get(ctx, key)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
 			}
 			return nil, fmt.Errorf("read subnet %s: %w", key, err)
@@ -362,30 +364,34 @@ func (s *RouteTableServiceImpl) publishSubnetEgressGateDecision(ctx context.Cont
 // destCidr, resolving each subnet's effective RT. Called on IGW Attach/Detach.
 // destCidr defaults to 0.0.0.0/0.
 func (s *RouteTableServiceImpl) PublishGateDecisionsForVPC(accountID, vpcID, destCidr string) {
+	s.publishGateDecisionsForVPC(context.Background(), accountID, vpcID, destCidr)
+}
+
+func (s *RouteTableServiceImpl) publishGateDecisionsForVPC(ctx context.Context, accountID, vpcID, destCidr string) {
 	if s.natsConn == nil || vpcID == "" {
 		return
 	}
 	if destCidr == "" {
 		destCidr = "0.0.0.0/0"
 	}
-	subnets, err := s.allSubnetsForVPC(accountID, vpcID)
+	subnets, err := s.allSubnetsForVPC(ctx, accountID, vpcID)
 	if err != nil {
-		slog.Warn("PublishGateDecisionsForVPC: enumerate subnets failed",
+		slog.WarnContext(ctx, "PublishGateDecisionsForVPC: enumerate subnets failed",
 			"vpcId", vpcID, "err", err)
 		return
 	}
 	for _, subnetID := range subnets {
-		s.publishSubnetEgressGateDecision(context.Background(), accountID, vpcID, subnetID, destCidr)
+		s.publishSubnetEgressGateDecision(ctx, accountID, vpcID, subnetID, destCidr)
 	}
 }
 
 // allSubnetsForVPC returns every SubnetId in vpcID by scanning the subnet KV.
 // Order is unspecified.
-func (s *RouteTableServiceImpl) allSubnetsForVPC(accountID, vpcID string) ([]string, error) {
+func (s *RouteTableServiceImpl) allSubnetsForVPC(ctx context.Context, accountID, vpcID string) ([]string, error) {
 	prefix := accountID + "."
-	keys, err := s.subnetKV.Keys()
+	keys, err := s.subnetKV.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, nil
 		}
 		return nil, fmt.Errorf("list subnet keys: %w", err)
@@ -395,9 +401,9 @@ func (s *RouteTableServiceImpl) allSubnetsForVPC(accountID, vpcID string) ([]str
 		if key == utils.VersionKey || !strings.HasPrefix(key, prefix) {
 			continue
 		}
-		entry, err := s.subnetKV.Get(key)
+		entry, err := s.subnetKV.Get(ctx, key)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
 			}
 			return nil, fmt.Errorf("read subnet %s: %w", key, err)
@@ -450,9 +456,9 @@ func (s *RouteTableServiceImpl) publishSubnetEgressGateDecisionForRT(ctx context
 // allRouteTablesForVPC returns all route tables belonging to a VPC.
 func (s *RouteTableServiceImpl) allRouteTablesForVPC(ctx context.Context, accountID, vpcID string) ([]RouteTableRecord, error) {
 	prefix := accountID + "."
-	keys, err := s.rtbKV.Keys()
+	keys, err := s.rtbKV.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, nil
 		}
 		return nil, errors.New(awserrors.ErrorServerInternal)
@@ -463,9 +469,9 @@ func (s *RouteTableServiceImpl) allRouteTablesForVPC(ctx context.Context, accoun
 		if key == utils.VersionKey || !strings.HasPrefix(key, prefix) {
 			continue
 		}
-		entry, err := s.rtbKV.Get(key)
+		entry, err := s.rtbKV.Get(ctx, key)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue // deleted between Keys() and Get()
 			}
 			slog.ErrorContext(ctx, "Failed to read route table during VPC scan", "key", key, "err", err)
@@ -581,7 +587,7 @@ func (s *RouteTableServiceImpl) DeleteRouteTable(ctx context.Context, input *ec2
 		}
 	}
 
-	if err := s.rtbKV.Delete(utils.AccountKey(accountID, rtbID)); err != nil {
+	if err := s.rtbKV.Delete(ctx, utils.AccountKey(accountID, rtbID)); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -619,8 +625,8 @@ func (s *RouteTableServiceImpl) DescribeRouteTables(ctx context.Context, input *
 	}
 
 	prefix := accountID + "."
-	keys, err := s.rtbKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	keys, err := s.rtbKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -632,9 +638,9 @@ func (s *RouteTableServiceImpl) DescribeRouteTables(ctx context.Context, input *
 			continue
 		}
 
-		entry, err := s.rtbKV.Get(key)
+		entry, err := s.rtbKV.Get(ctx, key)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
 			}
 			slog.ErrorContext(ctx, "Failed to read route table during describe", "key", key, "err", err)
@@ -701,7 +707,7 @@ func (s *RouteTableServiceImpl) CreateRoute(ctx context.Context, input *ec2.Crea
 	case input.GatewayId != nil && *input.GatewayId != "":
 		igwID := *input.GatewayId
 		// Verify IGW exists and is attached to the same VPC
-		igwEntry, err := s.igwKV.Get(utils.AccountKey(accountID, igwID))
+		igwEntry, err := s.igwKV.Get(ctx, utils.AccountKey(accountID, igwID))
 		if err != nil {
 			return nil, errors.New(awserrors.ErrorInvalidInternetGatewayIDNotFound)
 		}
@@ -729,7 +735,7 @@ func (s *RouteTableServiceImpl) CreateRoute(ctx context.Context, input *ec2.Crea
 	case input.NatGatewayId != nil && *input.NatGatewayId != "":
 		natgwID := *input.NatGatewayId
 		// Verify NAT GW exists and belongs to the same VPC
-		natgwEntry, err := s.natgwKV.Get(utils.AccountKey(accountID, natgwID))
+		natgwEntry, err := s.natgwKV.Get(ctx, utils.AccountKey(accountID, natgwID))
 		if err != nil {
 			return nil, errors.New(awserrors.ErrorInvalidNatGatewayIDNotFound)
 		}
@@ -878,7 +884,7 @@ func (s *RouteTableServiceImpl) ReplaceRoute(ctx context.Context, input *ec2.Rep
 	igwID := *input.GatewayId
 
 	// Verify IGW exists and is attached to same VPC
-	igwEntry, err := s.igwKV.Get(utils.AccountKey(accountID, igwID))
+	igwEntry, err := s.igwKV.Get(ctx, utils.AccountKey(accountID, igwID))
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidInternetGatewayIDNotFound)
 	}
@@ -921,7 +927,7 @@ func (s *RouteTableServiceImpl) AssociateRouteTable(ctx context.Context, input *
 	}
 
 	// Verify subnet exists and belongs to the same VPC
-	subnetEntry, err := s.subnetKV.Get(utils.AccountKey(accountID, subnetID))
+	subnetEntry, err := s.subnetKV.Get(ctx, utils.AccountKey(accountID, subnetID))
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidSubnetIDNotFound)
 	}
@@ -1005,9 +1011,9 @@ func (s *RouteTableServiceImpl) DisassociateRouteTable(ctx context.Context, inpu
 
 	// Search all route tables for this account to find the association
 	prefix := accountID + "."
-	keys, err := s.rtbKV.Keys()
+	keys, err := s.rtbKV.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, errors.New(awserrors.ErrorInvalidAssociationIDNotFound)
 		}
 		slog.ErrorContext(ctx, "Failed to list route table keys", "err", err)
@@ -1019,9 +1025,9 @@ func (s *RouteTableServiceImpl) DisassociateRouteTable(ctx context.Context, inpu
 			continue
 		}
 
-		entry, err := s.rtbKV.Get(key)
+		entry, err := s.rtbKV.Get(ctx, key)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
 			}
 			slog.ErrorContext(ctx, "Failed to read route table during disassociate scan", "key", key, "err", err)
@@ -1096,9 +1102,9 @@ func (s *RouteTableServiceImpl) ReplaceRouteTableAssociation(ctx context.Context
 
 	// Find and remove the association from its current route table
 	prefix := accountID + "."
-	keys, err := s.rtbKV.Keys()
+	keys, err := s.rtbKV.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, errors.New(awserrors.ErrorInvalidAssociationIDNotFound)
 		}
 		slog.ErrorContext(ctx, "Failed to list route table keys", "err", err)
@@ -1110,9 +1116,9 @@ func (s *RouteTableServiceImpl) ReplaceRouteTableAssociation(ctx context.Context
 			continue
 		}
 
-		entry, err := s.rtbKV.Get(key)
+		entry, err := s.rtbKV.Get(ctx, key)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
 			}
 			slog.ErrorContext(ctx, "Failed to read route table during replace scan", "key", key, "err", err)
@@ -1157,9 +1163,13 @@ func (s *RouteTableServiceImpl) ReplaceRouteTableAssociation(ctx context.Context
 				Main:          assoc.Main,
 			})
 			if err := s.putRouteTable(ctx, accountID, newRecord); err != nil {
-				// Compensate: restore association to old table to avoid data loss
+				// Restore the old association even when the request context caused the
+				// failed write, otherwise cancellation turns a rollback into a no-op.
 				oldRecord.Associations = append(oldRecord.Associations, assoc)
-				if restoreErr := s.putRouteTable(ctx, accountID, &oldRecord); restoreErr != nil {
+				restoreCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
+				restoreErr := s.putRouteTable(restoreCtx, accountID, &oldRecord)
+				cancel()
+				if restoreErr != nil {
 					slog.ErrorContext(ctx, "CRITICAL: ReplaceRouteTableAssociation partial failure, association lost",
 						"associationId", assocID, "oldRouteTableId", oldRecord.RouteTableId,
 						"newRouteTableId", newRtbID, "restoreErr", restoreErr, "originalErr", err)
@@ -1342,7 +1352,7 @@ func (s *RouteTableServiceImpl) publishNatGatewayDeleteEvents(ctx context.Contex
 	if s.natsConn == nil {
 		return
 	}
-	natgwEntry, err := s.natgwKV.Get(utils.AccountKey(accountID, natgwID))
+	natgwEntry, err := s.natgwKV.Get(ctx, utils.AccountKey(accountID, natgwID))
 	if err != nil {
 		slog.WarnContext(ctx, "NAT GW event: natgw lookup failed", "topic", "vpc.delete-nat-gateway", "natGatewayId", natgwID, "err", err)
 		return
@@ -1390,7 +1400,7 @@ func (s *RouteTableServiceImpl) publishNatGatewayEventsForAssociation(ctx contex
 		if r.NatGatewayId == "" {
 			continue
 		}
-		natgwEntry, err := s.natgwKV.Get(utils.AccountKey(accountID, r.NatGatewayId))
+		natgwEntry, err := s.natgwKV.Get(ctx, utils.AccountKey(accountID, r.NatGatewayId))
 		if err != nil {
 			slog.WarnContext(ctx, "NAT GW event: natgw lookup failed", "topic", topic, "natGatewayId", r.NatGatewayId, "err", err)
 			continue
@@ -1415,7 +1425,7 @@ func (s *RouteTableServiceImpl) publishNatGatewayEventForSubnet(ctx context.Cont
 	if s.natsConn == nil {
 		return
 	}
-	subnetEntry, err := s.subnetKV.Get(utils.AccountKey(accountID, subnetID))
+	subnetEntry, err := s.subnetKV.Get(ctx, utils.AccountKey(accountID, subnetID))
 	if err != nil {
 		slog.WarnContext(ctx, "NAT GW event: subnet lookup failed", "topic", topic, "subnetId", subnetID, "err", err)
 		return
@@ -1575,7 +1585,7 @@ func (s *RouteTableServiceImpl) ApplyRecordTags(input *ec2.CreateTagsInput, acco
 	if input == nil {
 		return nil
 	}
-	return utils.MirrorKVRecordTags(s.rtbKV, accountID, "rtb-", input.Resources,
+	return utils.MirrorKVRecordTags(context.Background(), s.rtbKV, accountID, "rtb-", input.Resources,
 		func(r *RouteTableRecord) *map[string]string { return &r.Tags },
 		utils.MergeTagsMut(input))
 }
@@ -1586,7 +1596,7 @@ func (s *RouteTableServiceImpl) RemoveRecordTags(input *ec2.DeleteTagsInput, acc
 	if input == nil {
 		return nil
 	}
-	return utils.MirrorKVRecordTags(s.rtbKV, accountID, "rtb-", input.Resources,
+	return utils.MirrorKVRecordTags(context.Background(), s.rtbKV, accountID, "rtb-", input.Resources,
 		func(r *RouteTableRecord) *map[string]string { return &r.Tags },
 		utils.RemoveTagsMut(input))
 }

spinifex/handlers/ec2/routetable/service_impl_test.go

@@ -1,7 +1,6 @@
 package handlers_ec2_routetable
 
 import (
-	"context"
 	"encoding/json"
 	"sync"
 	"testing"
@@ -16,6 +15,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/testutil"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -24,37 +24,49 @@ const testAccountID = "123456789012"
 
 func setupTestService(t *testing.T) *RouteTableServiceImpl {
 	t.Helper()
-	_, nc, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 
 	// Seed VPC KV
-	testutil.SeedKV(t, js, handlers_ec2_vpc.KVBucketVPCs, map[string][]byte{
+	seedKV(t, js, handlers_ec2_vpc.KVBucketVPCs, map[string][]byte{
 		utils.AccountKey(testAccountID, "vpc-test1"): []byte(`{"vpc_id":"vpc-test1","cidr_block":"10.0.0.0/16","state":"available"}`),
 	})
 
 	// Seed IGW KV (attached to vpc-test1)
-	testutil.SeedKV(t, js, handlers_ec2_igw.KVBucketIGW, map[string][]byte{
+	seedKV(t, js, handlers_ec2_igw.KVBucketIGW, map[string][]byte{
 		utils.AccountKey(testAccountID, "igw-test1"): []byte(`{"internet_gateway_id":"igw-test1","vpc_id":"vpc-test1","state":"available"}`),
 	})
 
 	// Seed subnet KV
-	testutil.SeedKV(t, js, handlers_ec2_vpc.KVBucketSubnets, map[string][]byte{
+	seedKV(t, js, handlers_ec2_vpc.KVBucketSubnets, map[string][]byte{
 		utils.AccountKey(testAccountID, "subnet-test1"): []byte(`{"subnet_id":"subnet-test1","vpc_id":"vpc-test1","cidr_block":"10.0.1.0/24","state":"available"}`),
 		utils.AccountKey(testAccountID, "subnet-priv1"): []byte(`{"subnet_id":"subnet-priv1","vpc_id":"vpc-test1","cidr_block":"10.0.11.0/24","state":"available"}`),
 	})
 
 	// Seed NAT Gateway KV
-	testutil.SeedKV(t, js, handlers_ec2_natgw.KVBucketNatGateways, map[string][]byte{
+	seedKV(t, js, handlers_ec2_natgw.KVBucketNatGateways, map[string][]byte{
 		utils.AccountKey(testAccountID, "nat-test1"): []byte(`{"nat_gateway_id":"nat-test1","vpc_id":"vpc-test1","subnet_id":"subnet-test1","public_ip":"192.168.1.100","state":"available"}`),
 	})
 
-	svc, err := NewRouteTableServiceImplWithNATS(nil, nc)
+	svc, err := NewRouteTableServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 	return svc
 }
 
+func seedKV(t *testing.T, js jetstream.JetStream, bucket string, entries map[string][]byte) jetstream.KeyValue {
+	t.Helper()
+	kv, err := js.CreateOrUpdateKeyValue(t.Context(), jetstream.KeyValueConfig{Bucket: bucket, History: 1})
+	require.NoError(t, err)
+	for key, value := range entries {
+		_, err := kv.Put(t.Context(), key, value)
+		require.NoError(t, err)
+	}
+	return kv
+}
+
 func createTestRtb(t *testing.T, svc *RouteTableServiceImpl) string {
 	t.Helper()
-	out, err := svc.CreateRouteTable(context.Background(), &ec2.CreateRouteTableInput{
+	out, err := svc.CreateRouteTable(t.Context(), &ec2.CreateRouteTableInput{
 		VpcId: aws.String("vpc-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
@@ -63,7 +75,7 @@ func createTestRtb(t *testing.T, svc *RouteTableServiceImpl) string {
 
 func TestCreateRouteTable(t *testing.T) {
 	svc := setupTestService(t)
-	out, err := svc.CreateRouteTable(context.Background(), &ec2.CreateRouteTableInput{
+	out, err := svc.CreateRouteTable(t.Context(), &ec2.CreateRouteTableInput{
 		VpcId: aws.String("vpc-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
@@ -86,7 +98,7 @@ func TestCreateRouteTable(t *testing.T) {
 // the NAT gateway (and its billable EIP) be reclaimed.
 func TestCreateRouteTable_PersistsTagsForTagFilterDiscovery(t *testing.T) {
 	svc := setupTestService(t)
-	_, err := svc.CreateRouteTable(context.Background(), &ec2.CreateRouteTableInput{
+	_, err := svc.CreateRouteTable(t.Context(), &ec2.CreateRouteTableInput{
 		VpcId: aws.String("vpc-test1"),
 		TagSpecifications: []*ec2.TagSpecification{{
 			ResourceType: aws.String(ec2.ResourceTypeRouteTable),
@@ -98,7 +110,7 @@ func TestCreateRouteTable_PersistsTagsForTagFilterDiscovery(t *testing.T) {
 	}, testAccountID)
 	require.NoError(t, err)
 
-	out, err := svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err := svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		Filters: []*ec2.Filter{
 			{Name: aws.String("tag:spinifex:eks-cluster"), Values: aws.StringSlice([]string{"alpha"})},
 			{Name: aws.String("tag:spinifex:eks-role"), Values: aws.StringSlice([]string{"cp-private-rt"})},
@@ -117,7 +129,7 @@ func TestCreateRouteTable_PersistsTagsForTagFilterDiscovery(t *testing.T) {
 
 func TestCreateRouteTable_VpcNotFound(t *testing.T) {
 	svc := setupTestService(t)
-	_, err := svc.CreateRouteTable(context.Background(), &ec2.CreateRouteTableInput{
+	_, err := svc.CreateRouteTable(t.Context(), &ec2.CreateRouteTableInput{
 		VpcId: aws.String("vpc-nonexistent"),
 	}, testAccountID)
 	assert.EqualError(t, err, awserrors.ErrorInvalidVpcIDNotFound)
@@ -127,13 +139,13 @@ func TestDeleteRouteTable(t *testing.T) {
 	svc := setupTestService(t)
 	rtbID := createTestRtb(t, svc)
 
-	_, err := svc.DeleteRouteTable(context.Background(), &ec2.DeleteRouteTableInput{
+	_, err := svc.DeleteRouteTable(t.Context(), &ec2.DeleteRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 	}, testAccountID)
 	require.NoError(t, err)
 
 	// Should be gone
-	_, err = svc.getRouteTable(context.Background(), testAccountID, rtbID)
+	_, err = svc.getRouteTable(t.Context(), testAccountID, rtbID)
 	assert.EqualError(t, err, awserrors.ErrorInvalidRouteTableIDNotFound)
 }
 
@@ -142,7 +154,7 @@ func TestDeleteRouteTable_Main(t *testing.T) {
 	record, err := svc.CreateRouteTableForVPC("vpc-test1", "10.0.0.0/16", testAccountID, true, "")
 	require.NoError(t, err)
 
-	_, err = svc.DeleteRouteTable(context.Background(), &ec2.DeleteRouteTableInput{
+	_, err = svc.DeleteRouteTable(t.Context(), &ec2.DeleteRouteTableInput{
 		RouteTableId: aws.String(record.RouteTableId),
 	}, testAccountID)
 	assert.EqualError(t, err, awserrors.ErrorDependencyViolation)
@@ -153,14 +165,14 @@ func TestDeleteRouteTable_WithAssociations(t *testing.T) {
 	rtbID := createTestRtb(t, svc)
 
 	// Associate a subnet
-	_, err := svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err := svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
 	// Should fail to delete
-	_, err = svc.DeleteRouteTable(context.Background(), &ec2.DeleteRouteTableInput{
+	_, err = svc.DeleteRouteTable(t.Context(), &ec2.DeleteRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 	}, testAccountID)
 	assert.EqualError(t, err, awserrors.ErrorDependencyViolation)
@@ -170,7 +182,7 @@ func TestDescribeRouteTables(t *testing.T) {
 	svc := setupTestService(t)
 	rtbID := createTestRtb(t, svc)
 
-	out, err := svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{}, testAccountID)
+	out, err := svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{}, testAccountID)
 	require.NoError(t, err)
 	assert.Len(t, out.RouteTables, 1)
 	assert.Equal(t, rtbID, *out.RouteTables[0].RouteTableId)
@@ -182,7 +194,7 @@ func TestDescribeRouteTables_FilterByVpcId(t *testing.T) {
 
 	name := "vpc-id"
 	val := "vpc-test1"
-	out, err := svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err := svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		Filters: []*ec2.Filter{{Name: &name, Values: []*string{&val}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -190,7 +202,7 @@ func TestDescribeRouteTables_FilterByVpcId(t *testing.T) {
 
 	// Filter by non-existent VPC
 	val2 := "vpc-nope"
-	out, err = svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err = svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		Filters: []*ec2.Filter{{Name: &name, Values: []*string{&val2}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -209,7 +221,7 @@ func TestDescribeRouteTables_FilterByMain(t *testing.T) {
 
 	name := "association.main"
 	val := "true"
-	out, err := svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err := svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		Filters: []*ec2.Filter{{Name: &name, Values: []*string{&val}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -222,14 +234,14 @@ func TestDescribeRouteTables_FilterByRouteState(t *testing.T) {
 	createTestRtb(t, svc) // has local route with state=active
 
 	// Exact match
-	out, err := svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err := svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		Filters: []*ec2.Filter{{Name: aws.String("route.state"), Values: []*string{aws.String("active")}}},
 	}, testAccountID)
 	require.NoError(t, err)
 	assert.Len(t, out.RouteTables, 1)
 
 	// Non-match
-	out, err = svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err = svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		Filters: []*ec2.Filter{{Name: aws.String("route.state"), Values: []*string{aws.String("blackhole")}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -240,13 +252,13 @@ func TestDescribeRouteTables_FilterByRouteOrigin(t *testing.T) {
 	svc := setupTestService(t)
 	createTestRtb(t, svc) // local route has origin=CreateRouteTable
 
-	out, err := svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err := svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		Filters: []*ec2.Filter{{Name: aws.String("route.origin"), Values: []*string{aws.String("CreateRouteTable")}}},
 	}, testAccountID)
 	require.NoError(t, err)
 	assert.Len(t, out.RouteTables, 1)
 
-	out, err = svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err = svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		Filters: []*ec2.Filter{{Name: aws.String("route.origin"), Values: []*string{aws.String("CreateRoute")}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -258,14 +270,14 @@ func TestDescribeRouteTables_FilterByRouteNatGatewayId(t *testing.T) {
 	createTestRtb(t, svc) // no NAT GW route
 
 	// No match — no routes have a nat-gateway-id
-	out, err := svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err := svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		Filters: []*ec2.Filter{{Name: aws.String("route.nat-gateway-id"), Values: []*string{aws.String("nat-000000")}}},
 	}, testAccountID)
 	require.NoError(t, err)
 	assert.Empty(t, out.RouteTables)
 
 	// Wildcard on empty field should not match
-	out, err = svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err = svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		Filters: []*ec2.Filter{{Name: aws.String("route.nat-gateway-id"), Values: []*string{aws.String("nat-*")}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -277,21 +289,21 @@ func TestDescribeRouteTables_FilterByOwnerId(t *testing.T) {
 	createTestRtb(t, svc)
 
 	// Exact match
-	out, err := svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err := svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		Filters: []*ec2.Filter{{Name: aws.String("owner-id"), Values: []*string{aws.String(testAccountID)}}},
 	}, testAccountID)
 	require.NoError(t, err)
 	assert.Len(t, out.RouteTables, 1)
 
 	// Non-match
-	out, err = svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err = svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		Filters: []*ec2.Filter{{Name: aws.String("owner-id"), Values: []*string{aws.String("999999999999")}}},
 	}, testAccountID)
 	require.NoError(t, err)
 	assert.Empty(t, out.RouteTables)
 
 	// Wildcard
-	out, err = svc.DescribeRouteTables(context.Background(), &ec2.DescribeRouteTablesInput{
+	out, err = svc.DescribeRouteTables(t.Context(), &ec2.DescribeRouteTablesInput{
 		Filters: []*ec2.Filter{{Name: aws.String("owner-id"), Values: []*string{aws.String("1234*")}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -302,7 +314,7 @@ func TestCreateRoute(t *testing.T) {
 	svc := setupTestService(t)
 	rtbID := createTestRtb(t, svc)
 
-	_, err := svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err := svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -310,7 +322,7 @@ func TestCreateRoute(t *testing.T) {
 	require.NoError(t, err)
 
 	// Verify route was added
-	record, err := svc.getRouteTable(context.Background(), testAccountID, rtbID)
+	record, err := svc.getRouteTable(t.Context(), testAccountID, rtbID)
 	require.NoError(t, err)
 	assert.Len(t, record.Routes, 2) // local + igw
 	assert.Equal(t, "0.0.0.0/0", record.Routes[1].DestinationCidrBlock)
@@ -321,7 +333,7 @@ func TestCreateRoute_DuplicateDestination(t *testing.T) {
 	svc := setupTestService(t)
 	rtbID := createTestRtb(t, svc)
 
-	_, err := svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err := svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -329,7 +341,7 @@ func TestCreateRoute_DuplicateDestination(t *testing.T) {
 	require.NoError(t, err)
 
 	// Duplicate should fail
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -341,20 +353,20 @@ func TestDeleteRoute(t *testing.T) {
 	svc := setupTestService(t)
 	rtbID := createTestRtb(t, svc)
 
-	_, err := svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err := svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	_, err = svc.DeleteRoute(context.Background(), &ec2.DeleteRouteInput{
+	_, err = svc.DeleteRoute(t.Context(), &ec2.DeleteRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	record, err := svc.getRouteTable(context.Background(), testAccountID, rtbID)
+	record, err := svc.getRouteTable(t.Context(), testAccountID, rtbID)
 	require.NoError(t, err)
 	assert.Len(t, record.Routes, 1) // only local remains
 }
@@ -363,7 +375,7 @@ func TestDeleteRoute_LocalRoute(t *testing.T) {
 	svc := setupTestService(t)
 	rtbID := createTestRtb(t, svc)
 
-	_, err := svc.DeleteRoute(context.Background(), &ec2.DeleteRouteInput{
+	_, err := svc.DeleteRoute(t.Context(), &ec2.DeleteRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("10.0.0.0/16"),
 	}, testAccountID)
@@ -374,7 +386,7 @@ func TestDeleteRoute_NotFound(t *testing.T) {
 	svc := setupTestService(t)
 	rtbID := createTestRtb(t, svc)
 
-	_, err := svc.DeleteRoute(context.Background(), &ec2.DeleteRouteInput{
+	_, err := svc.DeleteRoute(t.Context(), &ec2.DeleteRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("192.168.0.0/16"),
 	}, testAccountID)
@@ -385,7 +397,7 @@ func TestReplaceRoute(t *testing.T) {
 	svc := setupTestService(t)
 	rtbID := createTestRtb(t, svc)
 
-	_, err := svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err := svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -393,7 +405,7 @@ func TestReplaceRoute(t *testing.T) {
 	require.NoError(t, err)
 
 	// Replace target (same IGW for simplicity — validates the swap logic)
-	_, err = svc.ReplaceRoute(context.Background(), &ec2.ReplaceRouteInput{
+	_, err = svc.ReplaceRoute(t.Context(), &ec2.ReplaceRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -405,7 +417,7 @@ func TestAssociateRouteTable(t *testing.T) {
 	svc := setupTestService(t)
 	rtbID := createTestRtb(t, svc)
 
-	out, err := svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	out, err := svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-test1"),
 	}, testAccountID)
@@ -418,14 +430,14 @@ func TestAssociateRouteTable_DuplicateSubnet(t *testing.T) {
 	svc := setupTestService(t)
 	rtbID := createTestRtb(t, svc)
 
-	_, err := svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err := svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
 	// Second association should fail
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-test1"),
 	}, testAccountID)
@@ -449,7 +461,7 @@ func TestAssociateRouteTable_ConcurrentDistinctSubnets(t *testing.T) {
 		wg.Add(1)
 		go func(i int, sn string) {
 			defer wg.Done()
-			_, errs[i] = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+			_, errs[i] = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 				RouteTableId: aws.String(rtbID),
 				SubnetId:     aws.String(sn),
 			}, testAccountID)
@@ -461,7 +473,7 @@ func TestAssociateRouteTable_ConcurrentDistinctSubnets(t *testing.T) {
 		require.NoErrorf(t, err, "associate %s", subnets[i])
 	}
 
-	record, err := svc.getRouteTable(context.Background(), testAccountID, rtbID)
+	record, err := svc.getRouteTable(t.Context(), testAccountID, rtbID)
 	require.NoError(t, err)
 	got := make([]string, 0, len(record.Associations))
 	for _, a := range record.Associations {
@@ -474,19 +486,19 @@ func TestDisassociateRouteTable(t *testing.T) {
 	svc := setupTestService(t)
 	rtbID := createTestRtb(t, svc)
 
-	assocOut, err := svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	assocOut, err := svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	_, err = svc.DisassociateRouteTable(context.Background(), &ec2.DisassociateRouteTableInput{
+	_, err = svc.DisassociateRouteTable(t.Context(), &ec2.DisassociateRouteTableInput{
 		AssociationId: assocOut.AssociationId,
 	}, testAccountID)
 	require.NoError(t, err)
 
 	// Verify association removed
-	record, err := svc.getRouteTable(context.Background(), testAccountID, rtbID)
+	record, err := svc.getRouteTable(t.Context(), testAccountID, rtbID)
 	require.NoError(t, err)
 	assert.Empty(t, record.Associations)
 }
@@ -496,7 +508,7 @@ func TestDisassociateRouteTable_Main(t *testing.T) {
 	record, err := svc.CreateRouteTableForVPC("vpc-test1", "10.0.0.0/16", testAccountID, true, "")
 	require.NoError(t, err)
 
-	_, err = svc.DisassociateRouteTable(context.Background(), &ec2.DisassociateRouteTableInput{
+	_, err = svc.DisassociateRouteTable(t.Context(), &ec2.DisassociateRouteTableInput{
 		AssociationId: aws.String(record.Associations[0].AssociationId),
 	}, testAccountID)
 	assert.EqualError(t, err, awserrors.ErrorInvalidParameterValue)
@@ -507,13 +519,13 @@ func TestReplaceRouteTableAssociation(t *testing.T) {
 	rtb1ID := createTestRtb(t, svc)
 	rtb2ID := createTestRtb(t, svc)
 
-	assocOut, err := svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	assocOut, err := svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtb1ID),
 		SubnetId:     aws.String("subnet-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	replaceOut, err := svc.ReplaceRouteTableAssociation(context.Background(), &ec2.ReplaceRouteTableAssociationInput{
+	replaceOut, err := svc.ReplaceRouteTableAssociation(t.Context(), &ec2.ReplaceRouteTableAssociationInput{
 		AssociationId: assocOut.AssociationId,
 		RouteTableId:  aws.String(rtb2ID),
 	}, testAccountID)
@@ -522,12 +534,12 @@ func TestReplaceRouteTableAssociation(t *testing.T) {
 	assert.NotEqual(t, *assocOut.AssociationId, *replaceOut.NewAssociationId)
 
 	// Verify old table has no associations
-	oldRecord, err := svc.getRouteTable(context.Background(), testAccountID, rtb1ID)
+	oldRecord, err := svc.getRouteTable(t.Context(), testAccountID, rtb1ID)
 	require.NoError(t, err)
 	assert.Empty(t, oldRecord.Associations)
 
 	// Verify new table has the association
-	newRecord, err := svc.getRouteTable(context.Background(), testAccountID, rtb2ID)
+	newRecord, err := svc.getRouteTable(t.Context(), testAccountID, rtb2ID)
 	require.NoError(t, err)
 	require.Len(t, newRecord.Associations, 1)
 	assert.Equal(t, "subnet-test1", newRecord.Associations[0].SubnetId)
@@ -577,7 +589,7 @@ func TestAssociateRouteTable_PublishesNatGatewayEvent(t *testing.T) {
 	rtbID := createTestRtb(t, svc)
 
 	// Add NAT GW route BEFORE any subnet is associated (terraform ordering).
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		NatGatewayId:         aws.String("nat-test1"),
@@ -588,7 +600,7 @@ func TestAssociateRouteTable_PublishesNatGatewayEvent(t *testing.T) {
 	_, err = sub.NextMsg(100 * time.Millisecond)
 	assert.Error(t, err, "CreateRoute must not publish when table has no associations")
 
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
@@ -609,7 +621,7 @@ func TestAssociateRouteTable_NoNatGatewayRoute(t *testing.T) {
 
 	rtbID := createTestRtb(t, svc)
 
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
@@ -628,20 +640,20 @@ func TestDisassociateRouteTable_PublishesNatGatewayDeleteEvent(t *testing.T) {
 	defer func() { _ = delSub.Unsubscribe() }()
 
 	rtbID := createTestRtb(t, svc)
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		NatGatewayId:         aws.String("nat-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	assocOut, err := svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	assocOut, err := svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	_, err = svc.DisassociateRouteTable(context.Background(), &ec2.DisassociateRouteTableInput{
+	_, err = svc.DisassociateRouteTable(t.Context(), &ec2.DisassociateRouteTableInput{
 		AssociationId: assocOut.AssociationId,
 	}, testAccountID)
 	require.NoError(t, err)
@@ -662,12 +674,12 @@ func TestDeleteRoute_NATGW_PublishesDeleteForAssociatedSubnets(t *testing.T) {
 	defer func() { _ = delSub.Unsubscribe() }()
 
 	rtbID := createTestRtb(t, svc)
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
 	require.NoError(t, err)
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		NatGatewayId:         aws.String("nat-test1"),
@@ -677,7 +689,7 @@ func TestDeleteRoute_NATGW_PublishesDeleteForAssociatedSubnets(t *testing.T) {
 	natgwID, _ := receiveNatGWEvent(t, addSub, "10.0.11.0/24")
 	require.Equal(t, "nat-test1", natgwID)
 
-	_, err = svc.DeleteRoute(context.Background(), &ec2.DeleteRouteInput{
+	_, err = svc.DeleteRoute(t.Context(), &ec2.DeleteRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 	}, testAccountID)
@@ -702,14 +714,14 @@ func TestReplaceRouteTableAssociation_PublishesNatGatewayEvents(t *testing.T) {
 	oldRtb := createTestRtb(t, svc)
 	newRtb := createTestRtb(t, svc)
 
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(oldRtb),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		NatGatewayId:         aws.String("nat-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	assocOut, err := svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	assocOut, err := svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(oldRtb),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
@@ -718,7 +730,7 @@ func TestReplaceRouteTableAssociation_PublishesNatGatewayEvents(t *testing.T) {
 	_, err = addSub.NextMsg(time.Second)
 	require.NoError(t, err)
 
-	_, err = svc.ReplaceRouteTableAssociation(context.Background(), &ec2.ReplaceRouteTableAssociationInput{
+	_, err = svc.ReplaceRouteTableAssociation(t.Context(), &ec2.ReplaceRouteTableAssociationInput{
 		AssociationId: assocOut.AssociationId,
 		RouteTableId:  aws.String(newRtb),
 	}, testAccountID)
@@ -760,13 +772,13 @@ func TestCreateRoute_IGW_PublishesAddIGWRouteForAssociatedSubnets(t *testing.T)
 
 	// Associate subnet FIRST, then CreateRoute — exercises the "route after
 	// association" leg.
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -788,7 +800,7 @@ func TestAssociateRouteTable_PublishesAddIGWRouteForExistingRoute(t *testing.T)
 
 	rtbID := createTestRtb(t, svc)
 
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -799,7 +811,7 @@ func TestAssociateRouteTable_PublishesAddIGWRouteForExistingRoute(t *testing.T)
 	_, err = sub.NextMsg(100 * time.Millisecond)
 	assert.Error(t, err, "CreateRoute(IGW) must not publish when table has no associations")
 
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
@@ -819,20 +831,20 @@ func TestDeleteRoute_IGW_PublishesDeleteIGWRoute(t *testing.T) {
 	defer func() { _ = delSub.Unsubscribe() }()
 
 	rtbID := createTestRtb(t, svc)
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	_, err = svc.DeleteRoute(context.Background(), &ec2.DeleteRouteInput{
+	_, err = svc.DeleteRoute(t.Context(), &ec2.DeleteRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 	}, testAccountID)
@@ -852,20 +864,20 @@ func TestDisassociateRouteTable_PublishesDeleteIGWRoute(t *testing.T) {
 	defer func() { _ = delSub.Unsubscribe() }()
 
 	rtbID := createTestRtb(t, svc)
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	assocOut, err := svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	assocOut, err := svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	_, err = svc.DisassociateRouteTable(context.Background(), &ec2.DisassociateRouteTableInput{
+	_, err = svc.DisassociateRouteTable(t.Context(), &ec2.DisassociateRouteTableInput{
 		AssociationId: assocOut.AssociationId,
 	}, testAccountID)
 	require.NoError(t, err)
@@ -905,7 +917,7 @@ func TestCreateRoute_IGW_OnMainRT_FansOutToImplicitSubnets(t *testing.T) {
 	mainRT, err := svc.CreateRouteTableForVPC("vpc-test1", "10.0.0.0/16", testAccountID, true, "")
 	require.NoError(t, err)
 
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(mainRT.RouteTableId),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -931,13 +943,13 @@ func TestCreateRoute_IGW_OnMainRT_SkipsExplicitlyAssociatedSubnets(t *testing.T)
 	require.NoError(t, err)
 
 	customRtbID := createTestRtb(t, svc)
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(customRtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(mainRT.RouteTableId),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -959,7 +971,7 @@ func TestAssociateRouteTable_RemovesMainRTRoutesForJoiningSubnet(t *testing.T) {
 
 	mainRT, err := svc.CreateRouteTableForVPC("vpc-test1", "10.0.0.0/16", testAccountID, true, "")
 	require.NoError(t, err)
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(mainRT.RouteTableId),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -967,7 +979,7 @@ func TestAssociateRouteTable_RemovesMainRTRoutesForJoiningSubnet(t *testing.T) {
 	require.NoError(t, err)
 
 	customRtbID := createTestRtb(t, svc)
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(customRtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
@@ -1020,7 +1032,7 @@ func TestCreateRoute_IGW_PublishesUngateForAssociatedSubnets(t *testing.T) {
 	defer func() { _ = ungate.Unsubscribe() }()
 
 	rtbID := createTestRtb(t, svc)
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
@@ -1032,7 +1044,7 @@ func TestCreateRoute_IGW_PublishesUngateForAssociatedSubnets(t *testing.T) {
 		}
 	}
 
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -1054,12 +1066,12 @@ func TestDeleteRoute_IGW_PublishesGateForAssociatedSubnets(t *testing.T) {
 	defer func() { _ = gate.Unsubscribe() }()
 
 	rtbID := createTestRtb(t, svc)
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
 	require.NoError(t, err)
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -1071,7 +1083,7 @@ func TestDeleteRoute_IGW_PublishesGateForAssociatedSubnets(t *testing.T) {
 		}
 	}
 
-	_, err = svc.DeleteRoute(context.Background(), &ec2.DeleteRouteInput{
+	_, err = svc.DeleteRoute(t.Context(), &ec2.DeleteRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 	}, testAccountID)
@@ -1091,7 +1103,7 @@ func TestAssociateRouteTable_NoDefaultRoute_PublishesGate(t *testing.T) {
 	defer func() { _ = gate.Unsubscribe() }()
 
 	rtbID := createTestRtb(t, svc)
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
@@ -1112,14 +1124,14 @@ func TestAssociateRouteTable_HasIGWRoute_PublishesUngate(t *testing.T) {
 	defer func() { _ = ungate.Unsubscribe() }()
 
 	rtbID := createTestRtb(t, svc)
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
 
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
@@ -1142,13 +1154,13 @@ func TestDisassociateRouteTable_MainHasNoEgress_PublishesGate(t *testing.T) {
 	_, err = svc.CreateRouteTableForVPC("vpc-test1", "10.0.0.0/16", testAccountID, true, "")
 	require.NoError(t, err)
 	customRtbID := createTestRtb(t, svc)
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(customRtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
-	assocOut, err := svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	assocOut, err := svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(customRtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
@@ -1159,7 +1171,7 @@ func TestDisassociateRouteTable_MainHasNoEgress_PublishesGate(t *testing.T) {
 		}
 	}
 
-	_, err = svc.DisassociateRouteTable(context.Background(), &ec2.DisassociateRouteTableInput{
+	_, err = svc.DisassociateRouteTable(t.Context(), &ec2.DisassociateRouteTableInput{
 		AssociationId: assocOut.AssociationId,
 	}, testAccountID)
 	require.NoError(t, err)
@@ -1181,7 +1193,7 @@ func TestCreateRoute_NonDefaultPrefix_NoGateEvent(t *testing.T) {
 	defer func() { _ = ungate.Unsubscribe() }()
 
 	rtbID := createTestRtb(t, svc)
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
@@ -1193,7 +1205,7 @@ func TestCreateRoute_NonDefaultPrefix_NoGateEvent(t *testing.T) {
 		}
 	}
 
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("8.8.8.8/32"),
 		GatewayId:            aws.String("igw-test1"),
@@ -1216,7 +1228,7 @@ func TestCreateRoute_IGW_OnMainRT_FansOutUngateToImplicitSubnets(t *testing.T) {
 
 	mainRT, err := svc.CreateRouteTableForVPC("vpc-test1", "10.0.0.0/16", testAccountID, true, "")
 	require.NoError(t, err)
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(mainRT.RouteTableId),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -1236,7 +1248,7 @@ func TestDisassociateRouteTable_RestoresMainRTRoutesForDepartingSubnet(t *testin
 
 	mainRT, err := svc.CreateRouteTableForVPC("vpc-test1", "10.0.0.0/16", testAccountID, true, "")
 	require.NoError(t, err)
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(mainRT.RouteTableId),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
@@ -1244,7 +1256,7 @@ func TestDisassociateRouteTable_RestoresMainRTRoutesForDepartingSubnet(t *testin
 	require.NoError(t, err)
 
 	customRtbID := createTestRtb(t, svc)
-	assocOut, err := svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	assocOut, err := svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(customRtbID),
 		SubnetId:     aws.String("subnet-priv1"),
 	}, testAccountID)
@@ -1258,7 +1270,7 @@ func TestDisassociateRouteTable_RestoresMainRTRoutesForDepartingSubnet(t *testin
 		}
 	}
 
-	_, err = svc.DisassociateRouteTable(context.Background(), &ec2.DisassociateRouteTableInput{
+	_, err = svc.DisassociateRouteTable(t.Context(), &ec2.DisassociateRouteTableInput{
 		AssociationId: assocOut.AssociationId,
 	}, testAccountID)
 	require.NoError(t, err)
@@ -1298,13 +1310,13 @@ func TestPublishGateDecisionsForVPC_UngatesSubnetWithIGWRoute(t *testing.T) {
 
 	// Custom RT with 0.0.0.0/0 -> IGW, associated to subnet-test1.
 	rtbID := createTestRtb(t, svc)
-	_, err = svc.CreateRoute(context.Background(), &ec2.CreateRouteInput{
+	_, err = svc.CreateRoute(t.Context(), &ec2.CreateRouteInput{
 		RouteTableId:         aws.String(rtbID),
 		DestinationCidrBlock: aws.String("0.0.0.0/0"),
 		GatewayId:            aws.String("igw-test1"),
 	}, testAccountID)
 	require.NoError(t, err)
-	_, err = svc.AssociateRouteTable(context.Background(), &ec2.AssociateRouteTableInput{
+	_, err = svc.AssociateRouteTable(t.Context(), &ec2.AssociateRouteTableInput{
 		RouteTableId: aws.String(rtbID),
 		SubnetId:     aws.String("subnet-test1"),
 	}, testAccountID)
@@ -1353,15 +1365,15 @@ func TestEffectiveRouteTable_DeterministicWithDuplicateMain(t *testing.T) {
 		},
 		CreatedAt: now,
 	}
-	require.NoError(t, svc.putRouteTable(context.Background(), testAccountID, &orphan))
-	require.NoError(t, svc.putRouteTable(context.Background(), testAccountID, &real_))
+	require.NoError(t, svc.putRouteTable(t.Context(), testAccountID, &orphan))
+	require.NoError(t, svc.putRouteTable(t.Context(), testAccountID, &real_))
 
-	rt, err := svc.effectiveRouteTable(context.Background(), testAccountID, "vpc-test1", "subnet-priv1")
+	rt, err := svc.effectiveRouteTable(t.Context(), testAccountID, "vpc-test1", "subnet-priv1")
 	require.NoError(t, err)
 	require.NotNil(t, rt)
 	assert.Equal(t, "rtb-aaareal", rt.RouteTableId, "must prefer main RT with more routes")
 
-	main, err := svc.mainRouteTable(context.Background(), testAccountID, "vpc-test1")
+	main, err := svc.mainRouteTable(t.Context(), testAccountID, "vpc-test1")
 	require.NoError(t, err)
 	require.NotNil(t, main)
 	assert.Equal(t, "rtb-aaareal", main.RouteTableId, "mainRouteTable must use same tiebreak")

spinifex/handlers/ec2/snapshot/service_impl.go

@@ -22,12 +22,14 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/mulgadc/viperblock/viperblock"
 	vbs3 "github.com/mulgadc/viperblock/viperblock/backends/s3"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Ensure SnapshotServiceImpl implements SnapshotService.
@@ -36,6 +38,7 @@ var _ SnapshotService = (*SnapshotServiceImpl)(nil)
 const (
 	KVBucketVolumeSnapshots        = "spinifex-volume-snapshots"
 	KVBucketVolumeSnapshotsVersion = 1
+	snapshotCleanupTimeout         = 5 * time.Second
 )
 
 // SnapshotServiceImpl implements SnapshotService with S3-backed storage.
@@ -43,7 +46,7 @@ type SnapshotServiceImpl struct {
 	config   *config.Config
 	store    objectstore.ObjectStore
 	natsConn *nats.Conn
-	snapKV   nats.KeyValue
+	snapKV   jetstream.KeyValue
 	mutex    sync.RWMutex
 }
 
@@ -63,7 +66,7 @@ type SnapshotConfig struct {
 }
 
 // NewSnapshotServiceImplWithNATS creates a snapshot service with JetStream KV for volume-snapshot tracking.
-func NewSnapshotServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*SnapshotServiceImpl, nats.KeyValue, error) {
+func NewSnapshotServiceImplWithNATS(ctx context.Context, cfg *config.Config, natsConn *nats.Conn) (*SnapshotServiceImpl, jetstream.KeyValue, error) {
 	store := objectstore.NewS3ObjectStoreFromConfig(
 		cfg.Predastore.Host,
 		cfg.Predastore.Region,
@@ -71,16 +74,16 @@ func NewSnapshotServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*S
 		cfg.Predastore.SecretKey,
 	)
 
-	js, err := natsConn.JetStream()
+	js, err := jetstream.New(natsConn)
 	if err != nil {
-		return nil, nil, fmt.Errorf("failed to get JetStream context: %w", err)
+		return nil, nil, fmt.Errorf("create JetStream: %w", err)
 	}
 
-	kv, err := utils.GetOrCreateKVBucket(js, KVBucketVolumeSnapshots, 10)
+	kv, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketVolumeSnapshots, 10)
 	if err != nil {
-		return nil, nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketVolumeSnapshots, err)
+		return nil, nil, fmt.Errorf("create KV bucket %s: %w", KVBucketVolumeSnapshots, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketVolumeSnapshots, kv, KVBucketVolumeSnapshotsVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketVolumeSnapshots, kv, KVBucketVolumeSnapshotsVersion); err != nil {
 		return nil, nil, fmt.Errorf("migrate %s: %w", KVBucketVolumeSnapshots, err)
 	}
 
@@ -96,7 +99,7 @@ func NewSnapshotServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*S
 
 // NewSnapshotServiceImplWithStore creates a snapshot service with a custom ObjectStore (for testing).
 // An optional snapshotKV can be provided for KV-backed volume-snapshot tracking.
-func NewSnapshotServiceImplWithStore(cfg *config.Config, store objectstore.ObjectStore, natsConn *nats.Conn, snapshotKV ...nats.KeyValue) *SnapshotServiceImpl {
+func NewSnapshotServiceImplWithStore(cfg *config.Config, store objectstore.ObjectStore, natsConn *nats.Conn, snapshotKV ...jetstream.KeyValue) *SnapshotServiceImpl {
 	svc := &SnapshotServiceImpl{
 		config:   cfg,
 		store:    store,
@@ -271,14 +274,16 @@ func (s *SnapshotServiceImpl) CreateSnapshot(ctx context.Context, input *ec2.Cre
 
 	// Track the volume→snapshot dependency in KV before persisting to S3.
 	// This ensures we never have an untracked snapshot in S3.
-	if err := s.addSnapshotRef(volumeID, snapshotID); err != nil {
+	if err := s.addSnapshotRef(ctx, volumeID, snapshotID); err != nil {
 		slog.ErrorContext(ctx, "CreateSnapshot failed to add snapshot ref to KV", "snapshotId", snapshotID, "volumeId", volumeID, "err", err)
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
 	if err := s.putSnapshotConfig(snapshotID, snapshotCfg); err != nil {
 		slog.ErrorContext(ctx, "CreateSnapshot failed to write config", "snapshotId", snapshotID, "err", err)
-		_ = s.removeSnapshotRef(volumeID, snapshotID) // best-effort cleanup
+		if cleanupErr := s.removeSnapshotRefForCleanup(ctx, volumeID, snapshotID); cleanupErr != nil {
+			slog.Error("CreateSnapshot failed to roll back snapshot reference", "snapshotId", snapshotID, "volumeId", volumeID, "err", cleanupErr)
+		}
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -590,8 +595,8 @@ func (s *SnapshotServiceImpl) DeleteSnapshot(ctx context.Context, input *ec2.Del
 
 	// Remove from KV after S3 cleanup. Failure is logged but not fatal —
 	// a phantom entry safely blocks volume deletion rather than allowing it.
-	if err := s.removeSnapshotRef(cfg.VolumeID, snapshotID); err != nil {
-		slog.WarnContext(ctx, "DeleteSnapshot failed to remove snapshot ref from KV", "snapshotId", snapshotID, "volumeId", cfg.VolumeID, "err", err)
+	if err := s.removeSnapshotRefForCleanup(ctx, cfg.VolumeID, snapshotID); err != nil {
+		slog.Warn("DeleteSnapshot failed to remove snapshot ref from KV", "snapshotId", snapshotID, "volumeId", cfg.VolumeID, "err", err)
 	}
 
 	slog.InfoContext(ctx, "DeleteSnapshot completed", "snapshotId", snapshotID)
@@ -645,14 +650,16 @@ func (s *SnapshotServiceImpl) CopySnapshot(ctx context.Context, input *ec2.CopyS
 	maps.Copy(newCfg.Tags, sourceCfg.Tags)
 
 	// Track the volume→snapshot dependency in KV before persisting to S3.
-	if err := s.addSnapshotRef(sourceCfg.VolumeID, newSnapshotID); err != nil {
+	if err := s.addSnapshotRef(ctx, sourceCfg.VolumeID, newSnapshotID); err != nil {
 		slog.ErrorContext(ctx, "CopySnapshot failed to add snapshot ref to KV", "snapshotId", newSnapshotID, "volumeId", sourceCfg.VolumeID, "err", err)
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
 	if err := s.putSnapshotConfig(newSnapshotID, newCfg); err != nil {
 		slog.ErrorContext(ctx, "CopySnapshot failed to write config", "snapshotId", newSnapshotID, "err", err)
-		_ = s.removeSnapshotRef(sourceCfg.VolumeID, newSnapshotID) // best-effort cleanup
+		if cleanupErr := s.removeSnapshotRefForCleanup(ctx, sourceCfg.VolumeID, newSnapshotID); cleanupErr != nil {
+			slog.Error("CopySnapshot failed to roll back snapshot reference", "snapshotId", newSnapshotID, "volumeId", sourceCfg.VolumeID, "err", cleanupErr)
+		}
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -665,7 +672,7 @@ func (s *SnapshotServiceImpl) CopySnapshot(ctx context.Context, input *ec2.CopyS
 
 // addSnapshotRef adds snapshotID to the volume's snapshot list in KV.
 // Uses CAS (Create/Update with revision) to prevent lost updates under concurrency.
-func (s *SnapshotServiceImpl) addSnapshotRef(volumeID, snapshotID string) error {
+func (s *SnapshotServiceImpl) addSnapshotRef(ctx context.Context, volumeID, snapshotID string) error {
 	if s.snapKV == nil {
 		slog.Debug("addSnapshotRef: snapshotKV is nil, skipping", "volumeId", volumeID, "snapshotId", snapshotID)
 		return nil
@@ -673,11 +680,11 @@ func (s *SnapshotServiceImpl) addSnapshotRef(volumeID, snapshotID string) error
 
 	const maxRetries = 5
 	for attempt := range maxRetries {
-		entry, err := s.snapKV.Get(volumeID)
+		entry, err := s.snapKV.Get(ctx, volumeID)
 		var snapshots []string
 
 		if err != nil {
-			if !errors.Is(err, nats.ErrKeyNotFound) {
+			if !errors.Is(err, jetstream.ErrKeyNotFound) {
 				return fmt.Errorf("addSnapshotRef: failed to get KV key %s: %w", volumeID, err)
 			}
 			// Key doesn't exist yet — create with just this snapshot
@@ -685,7 +692,7 @@ func (s *SnapshotServiceImpl) addSnapshotRef(volumeID, snapshotID string) error
 			if err != nil {
 				return fmt.Errorf("addSnapshotRef: failed to marshal snapshot list: %w", err)
 			}
-			if _, err := s.snapKV.Create(volumeID, data); err != nil {
+			if _, err := s.snapKV.Create(ctx, volumeID, data); err != nil {
 				if attempt < maxRetries-1 {
 					continue // concurrent Create/Update — retry
 				}
@@ -706,7 +713,7 @@ func (s *SnapshotServiceImpl) addSnapshotRef(volumeID, snapshotID string) error
 			return fmt.Errorf("addSnapshotRef: failed to marshal snapshot list: %w", err)
 		}
 
-		if _, err := s.snapKV.Update(volumeID, data, entry.Revision()); err != nil {
+		if _, err := s.snapKV.Update(ctx, volumeID, data, entry.Revision()); err != nil {
 			if attempt < maxRetries-1 {
 				continue // concurrent update — retry
 			}
@@ -720,10 +727,18 @@ func (s *SnapshotServiceImpl) addSnapshotRef(volumeID, snapshotID string) error
 	return fmt.Errorf("addSnapshotRef: exhausted retries for KV key %s", volumeID)
 }
 
+// removeSnapshotRefForCleanup removes a reference with a bounded context that
+// survives request cancellation after the snapshot operation has committed.
+func (s *SnapshotServiceImpl) removeSnapshotRefForCleanup(ctx context.Context, volumeID, snapshotID string) error {
+	cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), snapshotCleanupTimeout)
+	defer cancel()
+	return s.removeSnapshotRef(cleanupCtx, volumeID, snapshotID)
+}
+
 // removeSnapshotRef removes snapshotID from the volume's snapshot list in KV.
 // Deletes the key if the list becomes empty.
 // Uses CAS (Update with revision) to prevent lost updates under concurrency.
-func (s *SnapshotServiceImpl) removeSnapshotRef(volumeID, snapshotID string) error {
+func (s *SnapshotServiceImpl) removeSnapshotRef(ctx context.Context, volumeID, snapshotID string) error {
 	if s.snapKV == nil {
 		slog.Debug("removeSnapshotRef: snapshotKV is nil, skipping", "volumeId", volumeID, "snapshotId", snapshotID)
 		return nil
@@ -731,9 +746,9 @@ func (s *SnapshotServiceImpl) removeSnapshotRef(volumeID, snapshotID string) err
 
 	const maxRetries = 5
 	for attempt := range maxRetries {
-		entry, err := s.snapKV.Get(volumeID)
+		entry, err := s.snapKV.Get(ctx, volumeID)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				return nil
 			}
 			return fmt.Errorf("removeSnapshotRef: failed to get KV key %s: %w", volumeID, err)
@@ -752,7 +767,7 @@ func (s *SnapshotServiceImpl) removeSnapshotRef(volumeID, snapshotID string) err
 		}
 
 		if len(filtered) == 0 {
-			if err := s.snapKV.Delete(volumeID); err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
+			if err := s.snapKV.Delete(ctx, volumeID); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
 				return fmt.Errorf("removeSnapshotRef: failed to delete KV key %s: %w", volumeID, err)
 			}
 		} else {
@@ -760,7 +775,7 @@ func (s *SnapshotServiceImpl) removeSnapshotRef(volumeID, snapshotID string) err
 			if err != nil {
 				return fmt.Errorf("removeSnapshotRef: failed to marshal snapshot list: %w", err)
 			}
-			if _, err := s.snapKV.Update(volumeID, data, entry.Revision()); err != nil {
+			if _, err := s.snapKV.Update(ctx, volumeID, data, entry.Revision()); err != nil {
 				if attempt < maxRetries-1 {
 					continue // concurrent update — retry
 				}
@@ -776,14 +791,14 @@ func (s *SnapshotServiceImpl) removeSnapshotRef(volumeID, snapshotID string) err
 }
 
 // volumeHasSnapshots returns true if the volume has any snapshots in KV.
-func (s *SnapshotServiceImpl) volumeHasSnapshots(volumeID string) (bool, error) {
+func (s *SnapshotServiceImpl) volumeHasSnapshots(ctx context.Context, volumeID string) (bool, error) {
 	if s.snapKV == nil {
 		return false, nil
 	}
 
-	entry, err := s.snapKV.Get(volumeID)
+	entry, err := s.snapKV.Get(ctx, volumeID)
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return false, nil
 		}
 		return false, fmt.Errorf("volumeHasSnapshots: failed to get KV key %s: %w", volumeID, err)

spinifex/handlers/ec2/snapshot/service_impl_test.go

@@ -8,7 +8,6 @@ import (
 	"net/http/httptest"
 	"strings"
 	"testing"
-	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
@@ -16,9 +15,9 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
+	"github.com/mulgadc/spinifex/spinifex/testutil"
 	"github.com/mulgadc/viperblock/viperblock"
-	"github.com/nats-io/nats-server/v2/server"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -588,30 +587,12 @@ func TestCopySnapshot_PreservesTags(t *testing.T) {
 }
 
 // setupTestNATSKV creates a NATS JetStream test server and returns a KV bucket for testing.
-func setupTestNATSKV(t *testing.T) nats.KeyValue {
+func setupTestNATSKV(t *testing.T) jetstream.KeyValue {
 	t.Helper()
-	opts := &server.Options{
-		Host:      "127.0.0.1",
-		Port:      -1,
-		JetStream: true,
-		StoreDir:  t.TempDir(),
-		NoLog:     true,
-		NoSigs:    true,
-	}
-	ns, err := server.NewServer(opts)
-	require.NoError(t, err)
-	go ns.Start()
-	require.True(t, ns.ReadyForConnections(5*time.Second))
-	t.Cleanup(func() { ns.Shutdown() })
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 
-	nc, err := nats.Connect(ns.ClientURL())
-	require.NoError(t, err)
-	t.Cleanup(func() { nc.Close() })
-
-	js, err := nc.JetStream()
-	require.NoError(t, err)
-
-	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{
+	kv, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{
 		Bucket: KVBucketVolumeSnapshots,
 	})
 	require.NoError(t, err)
@@ -622,10 +603,10 @@ func TestAddSnapshotRef(t *testing.T) {
 	kv := setupTestNATSKV(t)
 	svc := &SnapshotServiceImpl{snapKV: kv}
 
-	require.NoError(t, svc.addSnapshotRef("vol-1", "snap-a"))
-	require.NoError(t, svc.addSnapshotRef("vol-1", "snap-b"))
+	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a"))
+	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-b"))
 
-	entry, err := kv.Get("vol-1")
+	entry, err := kv.Get(t.Context(), "vol-1")
 	require.NoError(t, err)
 	var snapshots []string
 	require.NoError(t, json.Unmarshal(entry.Value(), &snapshots))
@@ -636,23 +617,23 @@ func TestRemoveSnapshotRef(t *testing.T) {
 	kv := setupTestNATSKV(t)
 	svc := &SnapshotServiceImpl{snapKV: kv}
 
-	require.NoError(t, svc.addSnapshotRef("vol-1", "snap-a"))
-	require.NoError(t, svc.addSnapshotRef("vol-1", "snap-b"))
+	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a"))
+	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-b"))
 
 	// Remove one
-	require.NoError(t, svc.removeSnapshotRef("vol-1", "snap-a"))
+	require.NoError(t, svc.removeSnapshotRef(t.Context(), "vol-1", "snap-a"))
 
-	entry, err := kv.Get("vol-1")
+	entry, err := kv.Get(t.Context(), "vol-1")
 	require.NoError(t, err)
 	var snapshots []string
 	require.NoError(t, json.Unmarshal(entry.Value(), &snapshots))
 	assert.Equal(t, []string{"snap-b"}, snapshots)
 
 	// Remove last — key should be deleted
-	require.NoError(t, svc.removeSnapshotRef("vol-1", "snap-b"))
+	require.NoError(t, svc.removeSnapshotRef(t.Context(), "vol-1", "snap-b"))
 
-	_, err = kv.Get("vol-1")
-	assert.ErrorIs(t, err, nats.ErrKeyNotFound)
+	_, err = kv.Get(t.Context(), "vol-1")
+	assert.ErrorIs(t, err, jetstream.ErrKeyNotFound)
 }
 
 func TestRemoveSnapshotRef_NonExistentKey(t *testing.T) {
@@ -660,7 +641,20 @@ func TestRemoveSnapshotRef_NonExistentKey(t *testing.T) {
 	svc := &SnapshotServiceImpl{snapKV: kv}
 
 	// Should not error on non-existent key
-	require.NoError(t, svc.removeSnapshotRef("vol-nonexistent", "snap-x"))
+	require.NoError(t, svc.removeSnapshotRef(t.Context(), "vol-nonexistent", "snap-x"))
+}
+
+func TestRemoveSnapshotRefForCleanupSurvivesCancellation(t *testing.T) {
+	kv := setupTestNATSKV(t)
+	svc := &SnapshotServiceImpl{snapKV: kv}
+	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a"))
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+	require.NoError(t, svc.removeSnapshotRefForCleanup(ctx, "vol-1", "snap-a"))
+
+	_, err := kv.Get(t.Context(), "vol-1")
+	require.ErrorIs(t, err, jetstream.ErrKeyNotFound)
 }
 
 func TestVolumeHasSnapshots(t *testing.T) {
@@ -668,19 +662,19 @@ func TestVolumeHasSnapshots(t *testing.T) {
 	svc := &SnapshotServiceImpl{snapKV: kv}
 
 	// No entry → false
-	has, err := svc.volumeHasSnapshots("vol-1")
+	has, err := svc.volumeHasSnapshots(t.Context(), "vol-1")
 	require.NoError(t, err)
 	assert.False(t, has)
 
 	// Add one → true
-	require.NoError(t, svc.addSnapshotRef("vol-1", "snap-a"))
-	has, err = svc.volumeHasSnapshots("vol-1")
+	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a"))
+	has, err = svc.volumeHasSnapshots(t.Context(), "vol-1")
 	require.NoError(t, err)
 	assert.True(t, has)
 
 	// Remove it → false
-	require.NoError(t, svc.removeSnapshotRef("vol-1", "snap-a"))
-	has, err = svc.volumeHasSnapshots("vol-1")
+	require.NoError(t, svc.removeSnapshotRef(t.Context(), "vol-1", "snap-a"))
+	has, err = svc.volumeHasSnapshots(t.Context(), "vol-1")
 	require.NoError(t, err)
 	assert.False(t, has)
 }
@@ -689,9 +683,9 @@ func TestKVNilFallback(t *testing.T) {
 	svc := &SnapshotServiceImpl{snapKV: nil}
 
 	// All methods should be no-ops when KV is nil
-	require.NoError(t, svc.addSnapshotRef("vol-1", "snap-a"))
-	require.NoError(t, svc.removeSnapshotRef("vol-1", "snap-a"))
-	has, err := svc.volumeHasSnapshots("vol-1")
+	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a"))
+	require.NoError(t, svc.removeSnapshotRef(t.Context(), "vol-1", "snap-a"))
+	has, err := svc.volumeHasSnapshots(t.Context(), "vol-1")
 	require.NoError(t, err)
 	assert.False(t, has)
 }
@@ -716,12 +710,12 @@ func TestCreateSnapshot_WritesKVEntry(t *testing.T) {
 	require.NoError(t, err)
 
 	// Verify KV entry exists
-	has, err := svc.volumeHasSnapshots(volumeID)
+	has, err := svc.volumeHasSnapshots(t.Context(), volumeID)
 	require.NoError(t, err)
 	assert.True(t, has)
 
 	// Verify snapshot ID is in the list
-	entry, err := kv.Get(volumeID)
+	entry, err := kv.Get(t.Context(), volumeID)
 	require.NoError(t, err)
 	var snapshots []string
 	require.NoError(t, json.Unmarshal(entry.Value(), &snapshots))
@@ -754,7 +748,7 @@ func TestDeleteSnapshot_RemovesKVEntry(t *testing.T) {
 	require.NoError(t, err)
 
 	// KV should now be empty for this volume
-	has, err := svc.volumeHasSnapshots(volumeID)
+	has, err := svc.volumeHasSnapshots(t.Context(), volumeID)
 	require.NoError(t, err)
 	assert.False(t, has)
 }
@@ -784,7 +778,7 @@ func TestCopySnapshot_AddsKVEntry(t *testing.T) {
 	require.NoError(t, err)
 
 	// Both snapshot IDs should be in KV
-	entry, err := kv.Get(volumeID)
+	entry, err := kv.Get(t.Context(), volumeID)
 	require.NoError(t, err)
 	var snapshots []string
 	require.NoError(t, json.Unmarshal(entry.Value(), &snapshots))

spinifex/handlers/ec2/spotinstance/service_impl.go

@@ -14,9 +14,11 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Ensure SpotInstanceServiceImpl implements SpotInstanceService.
@@ -28,30 +30,30 @@ var _ SpotInstanceService = (*SpotInstanceServiceImpl)(nil)
 type SpotInstanceServiceImpl struct {
 	config     *config.Config
 	natsConn   *nats.Conn
-	activeKV   nats.KeyValue
-	terminalKV nats.KeyValue
+	activeKV   jetstream.KeyValue
+	terminalKV jetstream.KeyValue
 }
 
 // NewSpotInstanceServiceImplWithNATS creates a spot instance service with NATS JetStream.
-func NewSpotInstanceServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*SpotInstanceServiceImpl, error) {
-	js, err := natsConn.JetStream()
+func NewSpotInstanceServiceImplWithNATS(ctx context.Context, cfg *config.Config, natsConn *nats.Conn) (*SpotInstanceServiceImpl, error) {
+	js, err := jetstream.New(natsConn)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get JetStream context: %w", err)
 	}
 
-	activeKV, err := utils.GetOrCreateKVBucket(js, KVBucketSpotRequests, 10)
+	activeKV, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketSpotRequests, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketSpotRequests, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketSpotRequests, activeKV, KVBucketSpotRequestsVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketSpotRequests, activeKV, KVBucketSpotRequestsVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketSpotRequests, err)
 	}
 
-	terminalKV, err := getOrCreateTerminalBucket(js)
+	terminalKV, err := getOrCreateTerminalBucket(ctx, js)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketSpotRequestsTerminal, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketSpotRequestsTerminal, terminalKV, KVBucketSpotRequestsVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketSpotRequestsTerminal, terminalKV, KVBucketSpotRequestsVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketSpotRequestsTerminal, err)
 	}
 
@@ -66,22 +68,27 @@ func NewSpotInstanceServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn)
 	}, nil
 }
 
-// getOrCreateTerminalBucket creates the terminal bucket with a TTL, mirroring
-// the terminated-instances bucket; GetOrCreateKVBucket cannot set a TTL.
-func getOrCreateTerminalBucket(js nats.JetStreamContext) (nats.KeyValue, error) {
-	kv, err := js.KeyValue(KVBucketSpotRequestsTerminal)
+// getOrCreateTerminalBucket opens the terminal bucket without changing its
+// existing replica or placement configuration, creating it only when absent.
+func getOrCreateTerminalBucket(ctx context.Context, js jetstream.KeyValueManager) (jetstream.KeyValue, error) {
+	kv, err := js.KeyValue(ctx, KVBucketSpotRequestsTerminal)
 	if err == nil {
 		return kv, nil
 	}
-	if !errors.Is(err, nats.ErrBucketNotFound) {
+	if !errors.Is(err, jetstream.ErrBucketNotFound) {
 		return nil, err
 	}
-	return js.CreateKeyValue(&nats.KeyValueConfig{
+
+	kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
 		Bucket:      KVBucketSpotRequestsTerminal,
 		Description: "Terminal Spot Instance Requests (auto-expire after 1 hour)",
 		History:     1,
 		TTL:         spotTerminalTTL,
 	})
+	if errors.Is(err, jetstream.ErrBucketExists) {
+		return js.KeyValue(ctx, KVBucketSpotRequestsTerminal)
+	}
+	return kv, err
 }
 
 // PutSpotInstanceRequests stores each request in the active bucket.
@@ -96,7 +103,7 @@ func (s *SpotInstanceServiceImpl) PutSpotInstanceRequests(ctx context.Context, i
 		if err != nil {
 			return nil, errors.New(awserrors.ErrorServerInternal)
 		}
-		if _, err := s.activeKV.Put(utils.AccountKey(accountID, sirID), data); err != nil {
+		if _, err := s.activeKV.Put(ctx, utils.AccountKey(accountID, sirID), data); err != nil {
 			slog.ErrorContext(ctx, "PutSpotInstanceRequests: KV put failed", "sirId", sirID, "err", err)
 			return nil, errors.New(awserrors.ErrorServerInternal)
 		}
@@ -133,11 +140,11 @@ func (s *SpotInstanceServiceImpl) DescribeSpotInstanceRequests(ctx context.Conte
 		}
 	}
 
-	active, err := s.listRequests(s.activeKV, accountID)
+	active, err := s.listRequests(ctx, s.activeKV, accountID)
 	if err != nil {
 		return nil, err
 	}
-	terminal, err := s.listRequests(s.terminalKV, accountID)
+	terminal, err := s.listRequests(ctx, s.terminalKV, accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -183,15 +190,15 @@ func (s *SpotInstanceServiceImpl) CancelSpotInstanceRequests(ctx context.Context
 		sirID := aws.StringValue(idPtr)
 		key := utils.AccountKey(accountID, sirID)
 
-		record, err := s.readRecord(s.activeKV, key)
+		record, err := s.readRecord(ctx, s.activeKV, key)
 		if err == nil {
 			record.Request.State = aws.String(ec2.SpotInstanceStateCancelled)
 			setSpotStatus(record.Request, SpotStatusCodeCanceledInstanceRunning,
 				"Spot Instance request cancelled; the instance is still running.")
-			if err := s.moveToTerminal(key, record); err != nil {
+			if err := s.moveToTerminal(ctx, key, record); err != nil {
 				return nil, err
 			}
-		} else if !errors.Is(err, nats.ErrKeyNotFound) {
+		} else if !errors.Is(err, jetstream.ErrKeyNotFound) {
 			return nil, errors.New(awserrors.ErrorServerInternal)
 		}
 
@@ -213,9 +220,9 @@ func (s *SpotInstanceServiceImpl) CloseForInstance(ctx context.Context, instance
 		return nil
 	}
 
-	keys, err := s.activeKV.Keys()
+	keys, err := s.activeKV.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil
 		}
 		return errors.New(awserrors.ErrorServerInternal)
@@ -226,7 +233,7 @@ func (s *SpotInstanceServiceImpl) CloseForInstance(ctx context.Context, instance
 		if key == utils.VersionKey || !strings.HasPrefix(key, prefix) {
 			continue
 		}
-		record, err := s.readRecord(s.activeKV, key)
+		record, err := s.readRecord(ctx, s.activeKV, key)
 		if err != nil {
 			continue
 		}
@@ -237,7 +244,7 @@ func (s *SpotInstanceServiceImpl) CloseForInstance(ctx context.Context, instance
 		record.Request.State = aws.String(ec2.SpotInstanceStateClosed)
 		setSpotStatus(record.Request, SpotStatusCodeInstanceTerminatedByUser,
 			"Spot Instance request closed; the instance was terminated by the user.")
-		if err := s.moveToTerminal(key, record); err != nil {
+		if err := s.moveToTerminal(ctx, key, record); err != nil {
 			return err
 		}
 		slog.InfoContext(ctx, "CloseForInstance moved SIR to terminal", "instanceId", instanceID,
@@ -249,10 +256,10 @@ func (s *SpotInstanceServiceImpl) CloseForInstance(ctx context.Context, instance
 }
 
 // listRequests reads all account-scoped requests from a bucket.
-func (s *SpotInstanceServiceImpl) listRequests(kv nats.KeyValue, accountID string) ([]*ec2.SpotInstanceRequest, error) {
-	keys, err := kv.Keys()
+func (s *SpotInstanceServiceImpl) listRequests(ctx context.Context, kv jetstream.KeyValue, accountID string) ([]*ec2.SpotInstanceRequest, error) {
+	keys, err := kv.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, nil
 		}
 		return nil, errors.New(awserrors.ErrorServerInternal)
@@ -264,7 +271,7 @@ func (s *SpotInstanceServiceImpl) listRequests(kv nats.KeyValue, accountID strin
 		if key == utils.VersionKey || !strings.HasPrefix(key, prefix) {
 			continue
 		}
-		record, err := s.readRecord(kv, key)
+		record, err := s.readRecord(ctx, kv, key)
 		if err != nil {
 			slog.Warn("Failed to read spot request record", "key", key, "err", err)
 			continue
@@ -275,9 +282,9 @@ func (s *SpotInstanceServiceImpl) listRequests(kv nats.KeyValue, accountID strin
 }
 
 // readRecord fetches and unmarshals a single record. The KV Get error (e.g.
-// nats.ErrKeyNotFound) is returned unwrapped so callers can match on it.
-func (s *SpotInstanceServiceImpl) readRecord(kv nats.KeyValue, key string) (*SpotRequestRecord, error) {
-	entry, err := kv.Get(key)
+// jetstream.ErrKeyNotFound) is returned unwrapped so callers can match on it.
+func (s *SpotInstanceServiceImpl) readRecord(ctx context.Context, kv jetstream.KeyValue, key string) (*SpotRequestRecord, error) {
+	entry, err := kv.Get(ctx, key)
 	if err != nil {
 		return nil, err
 	}
@@ -289,16 +296,16 @@ func (s *SpotInstanceServiceImpl) readRecord(kv nats.KeyValue, key string) (*Spo
 }
 
 // moveToTerminal writes the record to the terminal bucket and deletes it from active.
-func (s *SpotInstanceServiceImpl) moveToTerminal(key string, record *SpotRequestRecord) error {
+func (s *SpotInstanceServiceImpl) moveToTerminal(ctx context.Context, key string, record *SpotRequestRecord) error {
 	data, err := json.Marshal(record)
 	if err != nil {
 		return errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.terminalKV.Put(key, data); err != nil {
+	if _, err := s.terminalKV.Put(ctx, key, data); err != nil {
 		slog.Error("moveToTerminal: terminal put failed", "key", key, "err", err)
 		return errors.New(awserrors.ErrorServerInternal)
 	}
-	if err := s.activeKV.Delete(key); err != nil {
+	if err := s.activeKV.Delete(ctx, key); err != nil {
 		slog.Error("moveToTerminal: active delete failed", "key", key, "err", err)
 		return errors.New(awserrors.ErrorServerInternal)
 	}

spinifex/handlers/ec2/spotinstance/service_impl_test.go

@@ -1,13 +1,14 @@
 package handlers_ec2_spotinstance
 
 import (
-	"context"
 	"testing"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/testutil"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -18,11 +19,43 @@ func setupTestService(t *testing.T) *SpotInstanceServiceImpl {
 	t.Helper()
 	_, nc, _ := testutil.StartTestJetStream(t)
 
-	svc, err := NewSpotInstanceServiceImplWithNATS(nil, nc)
+	svc, err := NewSpotInstanceServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 	return svc
 }
 
+func TestTerminalBucketTTL(t *testing.T) {
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+
+	_, err := NewSpotInstanceServiceImplWithNATS(t.Context(), nil, nc)
+	require.NoError(t, err)
+	kv, err := js.KeyValue(t.Context(), KVBucketSpotRequestsTerminal)
+	require.NoError(t, err)
+	status, err := kv.Status(t.Context())
+	require.NoError(t, err)
+	assert.Equal(t, spotTerminalTTL, status.TTL())
+}
+
+func TestGetOrCreateTerminalBucketPreservesExistingConfiguration(t *testing.T) {
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	const existingTTL = 2 * time.Hour
+
+	_, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{
+		Bucket:  KVBucketSpotRequestsTerminal,
+		History: 1,
+		TTL:     existingTTL,
+	})
+	require.NoError(t, err)
+
+	kv, err := getOrCreateTerminalBucket(t.Context(), js)
+	require.NoError(t, err)
+	status, err := kv.Status(t.Context())
+	require.NoError(t, err)
+	require.Equal(t, existingTTL, status.TTL())
+}
+
 // makeSIR builds an active/fulfilled Spot Instance Request for tests.
 func makeSIR(sirID, instanceID string) *ec2.SpotInstanceRequest {
 	return &ec2.SpotInstanceRequest{
@@ -42,13 +75,13 @@ func makeSIR(sirID, instanceID string) *ec2.SpotInstanceRequest {
 
 func putSIRs(t *testing.T, svc *SpotInstanceServiceImpl, reqs ...*ec2.SpotInstanceRequest) {
 	t.Helper()
-	_, err := svc.PutSpotInstanceRequests(context.Background(), &PutSpotRequestsInput{Requests: reqs}, testAccountID)
+	_, err := svc.PutSpotInstanceRequests(t.Context(), &PutSpotRequestsInput{Requests: reqs}, testAccountID)
 	require.NoError(t, err)
 }
 
 func describeAll(t *testing.T, svc *SpotInstanceServiceImpl, accountID string) []*ec2.SpotInstanceRequest {
 	t.Helper()
-	out, err := svc.DescribeSpotInstanceRequests(context.Background(), &ec2.DescribeSpotInstanceRequestsInput{}, accountID)
+	out, err := svc.DescribeSpotInstanceRequests(t.Context(), &ec2.DescribeSpotInstanceRequestsInput{}, accountID)
 	require.NoError(t, err)
 	return out.SpotInstanceRequests
 }
@@ -73,7 +106,7 @@ func TestPutDescribe_RoundTrip(t *testing.T) {
 
 func TestPutSpotInstanceRequests_MissingID(t *testing.T) {
 	svc := setupTestService(t)
-	_, err := svc.PutSpotInstanceRequests(context.Background(), &PutSpotRequestsInput{
+	_, err := svc.PutSpotInstanceRequests(t.Context(), &PutSpotRequestsInput{
 		Requests: []*ec2.SpotInstanceRequest{{InstanceId: aws.String("i-noid")}},
 	}, testAccountID)
 	require.Error(t, err)
@@ -99,7 +132,7 @@ func TestDescribe_FilterByID(t *testing.T) {
 	svc := setupTestService(t)
 	putSIRs(t, svc, makeSIR("sir-aaa", "i-aaa"), makeSIR("sir-bbb", "i-bbb"))
 
-	out, err := svc.DescribeSpotInstanceRequests(context.Background(), &ec2.DescribeSpotInstanceRequestsInput{
+	out, err := svc.DescribeSpotInstanceRequests(t.Context(), &ec2.DescribeSpotInstanceRequestsInput{
 		SpotInstanceRequestIds: []*string{aws.String("sir-aaa")},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -111,7 +144,7 @@ func TestDescribe_UnknownID(t *testing.T) {
 	svc := setupTestService(t)
 	putSIRs(t, svc, makeSIR("sir-aaa", "i-aaa"))
 
-	_, err := svc.DescribeSpotInstanceRequests(context.Background(), &ec2.DescribeSpotInstanceRequestsInput{
+	_, err := svc.DescribeSpotInstanceRequests(t.Context(), &ec2.DescribeSpotInstanceRequestsInput{
 		SpotInstanceRequestIds: []*string{aws.String("sir-ghost")},
 	}, testAccountID)
 	require.Error(t, err)
@@ -123,12 +156,12 @@ func TestDescribe_FilterByState(t *testing.T) {
 	putSIRs(t, svc, makeSIR("sir-aaa", "i-aaa"), makeSIR("sir-bbb", "i-bbb"))
 
 	// Cancel one so it becomes terminal/cancelled.
-	_, err := svc.CancelSpotInstanceRequests(context.Background(), &ec2.CancelSpotInstanceRequestsInput{
+	_, err := svc.CancelSpotInstanceRequests(t.Context(), &ec2.CancelSpotInstanceRequestsInput{
 		SpotInstanceRequestIds: []*string{aws.String("sir-bbb")},
 	}, testAccountID)
 	require.NoError(t, err)
 
-	out, err := svc.DescribeSpotInstanceRequests(context.Background(), &ec2.DescribeSpotInstanceRequestsInput{
+	out, err := svc.DescribeSpotInstanceRequests(t.Context(), &ec2.DescribeSpotInstanceRequestsInput{
 		Filters: []*ec2.Filter{{Name: aws.String("state"), Values: []*string{aws.String("active")}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -140,7 +173,7 @@ func TestDescribe_FilterByInstanceID(t *testing.T) {
 	svc := setupTestService(t)
 	putSIRs(t, svc, makeSIR("sir-aaa", "i-aaa"), makeSIR("sir-bbb", "i-bbb"))
 
-	out, err := svc.DescribeSpotInstanceRequests(context.Background(), &ec2.DescribeSpotInstanceRequestsInput{
+	out, err := svc.DescribeSpotInstanceRequests(t.Context(), &ec2.DescribeSpotInstanceRequestsInput{
 		Filters: []*ec2.Filter{{Name: aws.String("instance-id"), Values: []*string{aws.String("i-bbb")}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -154,7 +187,7 @@ func TestDescribe_FilterByLaunchImageID(t *testing.T) {
 	other.LaunchSpecification.ImageId = aws.String("ami-99999999")
 	putSIRs(t, svc, makeSIR("sir-aaa", "i-aaa"), other)
 
-	out, err := svc.DescribeSpotInstanceRequests(context.Background(), &ec2.DescribeSpotInstanceRequestsInput{
+	out, err := svc.DescribeSpotInstanceRequests(t.Context(), &ec2.DescribeSpotInstanceRequestsInput{
 		Filters: []*ec2.Filter{{Name: aws.String("launch.image-id"), Values: []*string{aws.String("ami-12345678")}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -168,7 +201,7 @@ func TestDescribe_FilterByTag(t *testing.T) {
 	tagged.Tags = []*ec2.Tag{{Key: aws.String("env"), Value: aws.String("prod")}}
 	putSIRs(t, svc, tagged, makeSIR("sir-bbb", "i-bbb"))
 
-	out, err := svc.DescribeSpotInstanceRequests(context.Background(), &ec2.DescribeSpotInstanceRequestsInput{
+	out, err := svc.DescribeSpotInstanceRequests(t.Context(), &ec2.DescribeSpotInstanceRequestsInput{
 		Filters: []*ec2.Filter{{Name: aws.String("tag:env"), Values: []*string{aws.String("prod")}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -176,7 +209,7 @@ func TestDescribe_FilterByTag(t *testing.T) {
 	assert.Equal(t, "sir-aaa", aws.StringValue(out.SpotInstanceRequests[0].SpotInstanceRequestId))
 
 	// tag-key matches any value for the key.
-	out, err = svc.DescribeSpotInstanceRequests(context.Background(), &ec2.DescribeSpotInstanceRequestsInput{
+	out, err = svc.DescribeSpotInstanceRequests(t.Context(), &ec2.DescribeSpotInstanceRequestsInput{
 		Filters: []*ec2.Filter{{Name: aws.String("tag-key"), Values: []*string{aws.String("env")}}},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -186,7 +219,7 @@ func TestDescribe_FilterByTag(t *testing.T) {
 
 func TestDescribe_InvalidFilter(t *testing.T) {
 	svc := setupTestService(t)
-	_, err := svc.DescribeSpotInstanceRequests(context.Background(), &ec2.DescribeSpotInstanceRequestsInput{
+	_, err := svc.DescribeSpotInstanceRequests(t.Context(), &ec2.DescribeSpotInstanceRequestsInput{
 		Filters: []*ec2.Filter{{Name: aws.String("bogus-filter"), Values: []*string{aws.String("x")}}},
 	}, testAccountID)
 	require.Error(t, err)
@@ -199,7 +232,7 @@ func TestCancel_MovesActiveToTerminal(t *testing.T) {
 	svc := setupTestService(t)
 	putSIRs(t, svc, makeSIR("sir-aaa", "i-aaa"))
 
-	out, err := svc.CancelSpotInstanceRequests(context.Background(), &ec2.CancelSpotInstanceRequestsInput{
+	out, err := svc.CancelSpotInstanceRequests(t.Context(), &ec2.CancelSpotInstanceRequestsInput{
 		SpotInstanceRequestIds: []*string{aws.String("sir-aaa")},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -219,11 +252,11 @@ func TestCancel_Idempotent(t *testing.T) {
 	putSIRs(t, svc, makeSIR("sir-aaa", "i-aaa"))
 
 	in := &ec2.CancelSpotInstanceRequestsInput{SpotInstanceRequestIds: []*string{aws.String("sir-aaa")}}
-	_, err := svc.CancelSpotInstanceRequests(context.Background(), in, testAccountID)
+	_, err := svc.CancelSpotInstanceRequests(t.Context(), in, testAccountID)
 	require.NoError(t, err)
 
 	// Second cancel of an already-terminal request succeeds without duplicating.
-	out, err := svc.CancelSpotInstanceRequests(context.Background(), in, testAccountID)
+	out, err := svc.CancelSpotInstanceRequests(t.Context(), in, testAccountID)
 	require.NoError(t, err)
 	require.Len(t, out.CancelledSpotInstanceRequests, 1)
 
@@ -234,7 +267,7 @@ func TestCancel_Idempotent(t *testing.T) {
 
 func TestCancel_AbsentIDIdempotent(t *testing.T) {
 	svc := setupTestService(t)
-	out, err := svc.CancelSpotInstanceRequests(context.Background(), &ec2.CancelSpotInstanceRequestsInput{
+	out, err := svc.CancelSpotInstanceRequests(t.Context(), &ec2.CancelSpotInstanceRequestsInput{
 		SpotInstanceRequestIds: []*string{aws.String("sir-ghost")},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -244,7 +277,7 @@ func TestCancel_AbsentIDIdempotent(t *testing.T) {
 
 func TestCancel_MissingIDs(t *testing.T) {
 	svc := setupTestService(t)
-	_, err := svc.CancelSpotInstanceRequests(context.Background(), &ec2.CancelSpotInstanceRequestsInput{}, testAccountID)
+	_, err := svc.CancelSpotInstanceRequests(t.Context(), &ec2.CancelSpotInstanceRequestsInput{}, testAccountID)
 	require.Error(t, err)
 	assert.Equal(t, awserrors.ErrorMissingParameter, err.Error())
 }
@@ -255,9 +288,9 @@ func TestCloseForInstance_MovesMatching(t *testing.T) {
 	svc := setupTestService(t)
 	putSIRs(t, svc, makeSIR("sir-aaa", "i-aaa"), makeSIR("sir-bbb", "i-bbb"))
 
-	require.NoError(t, svc.CloseForInstance(context.Background(), "i-bbb", testAccountID))
+	require.NoError(t, svc.CloseForInstance(t.Context(), "i-bbb", testAccountID))
 
-	out, err := svc.DescribeSpotInstanceRequests(context.Background(), &ec2.DescribeSpotInstanceRequestsInput{
+	out, err := svc.DescribeSpotInstanceRequests(t.Context(), &ec2.DescribeSpotInstanceRequestsInput{
 		SpotInstanceRequestIds: []*string{aws.String("sir-bbb")},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -266,7 +299,7 @@ func TestCloseForInstance_MovesMatching(t *testing.T) {
 	assert.Equal(t, SpotStatusCodeInstanceTerminatedByUser, aws.StringValue(out.SpotInstanceRequests[0].Status.Code))
 
 	// sir-aaa is untouched and still active.
-	out, err = svc.DescribeSpotInstanceRequests(context.Background(), &ec2.DescribeSpotInstanceRequestsInput{
+	out, err = svc.DescribeSpotInstanceRequests(t.Context(), &ec2.DescribeSpotInstanceRequestsInput{
 		SpotInstanceRequestIds: []*string{aws.String("sir-aaa")},
 	}, testAccountID)
 	require.NoError(t, err)
@@ -277,7 +310,7 @@ func TestCloseForInstance_NoMatchNoOp(t *testing.T) {
 	svc := setupTestService(t)
 	putSIRs(t, svc, makeSIR("sir-aaa", "i-aaa"))
 
-	require.NoError(t, svc.CloseForInstance(context.Background(), "i-nomatch", testAccountID))
+	require.NoError(t, svc.CloseForInstance(t.Context(), "i-nomatch", testAccountID))
 
 	got := describeAll(t, svc, testAccountID)
 	require.Len(t, got, 1)
@@ -286,7 +319,7 @@ func TestCloseForInstance_NoMatchNoOp(t *testing.T) {
 
 func TestCloseForInstance_EmptyInstanceID(t *testing.T) {
 	svc := setupTestService(t)
-	require.NoError(t, svc.CloseForInstance(context.Background(), "", testAccountID))
+	require.NoError(t, svc.CloseForInstance(t.Context(), "", testAccountID))
 }
 
 // --- Merge across buckets ---
@@ -295,11 +328,11 @@ func TestDescribe_MergesBothBuckets(t *testing.T) {
 	svc := setupTestService(t)
 	putSIRs(t, svc, makeSIR("sir-active", "i-1"), makeSIR("sir-cancel", "i-2"), makeSIR("sir-close", "i-3"))
 
-	_, err := svc.CancelSpotInstanceRequests(context.Background(), &ec2.CancelSpotInstanceRequestsInput{
+	_, err := svc.CancelSpotInstanceRequests(t.Context(), &ec2.CancelSpotInstanceRequestsInput{
 		SpotInstanceRequestIds: []*string{aws.String("sir-cancel")},
 	}, testAccountID)
 	require.NoError(t, err)
-	require.NoError(t, svc.CloseForInstance(context.Background(), "i-3", testAccountID))
+	require.NoError(t, svc.CloseForInstance(t.Context(), "i-3", testAccountID))
 
 	states := map[string]string{}
 	for _, r := range describeAll(t, svc, testAccountID) {

spinifex/handlers/ec2/volume/service_impl.go

@@ -29,6 +29,7 @@ import (
 	"github.com/mulgadc/viperblock/viperblock"
 	s3backend "github.com/mulgadc/viperblock/viperblock/backends/s3"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 const (
@@ -57,13 +58,13 @@ type VolumeServiceImpl struct {
 	store      objectstore.ObjectStore
 	bucketName string
 	natsConn   *nats.Conn
-	snapshotKV nats.KeyValue
+	snapshotKV jetstream.KeyValue
 }
 
 // NewVolumeServiceImpl creates a new daemon-side volume service.
 // snapshotKV is optional — when non-nil, DeleteVolume uses O(1) KV lookup
 // instead of scanning all snapshots in S3.
-func NewVolumeServiceImpl(cfg *config.Config, natsConn *nats.Conn, snapshotKV nats.KeyValue) *VolumeServiceImpl {
+func NewVolumeServiceImpl(cfg *config.Config, natsConn *nats.Conn, snapshotKV jetstream.KeyValue) *VolumeServiceImpl {
 	store := objectstore.NewS3ObjectStoreFromConfig(
 		cfg.Predastore.Host,
 		cfg.Predastore.Region,
@@ -81,7 +82,7 @@ func NewVolumeServiceImpl(cfg *config.Config, natsConn *nats.Conn, snapshotKV na
 }
 
 // NewVolumeServiceImplWithStore creates a volume service with a custom ObjectStore (for testing).
-func NewVolumeServiceImplWithStore(cfg *config.Config, store objectstore.ObjectStore, natsConn *nats.Conn, snapshotKV ...nats.KeyValue) *VolumeServiceImpl {
+func NewVolumeServiceImplWithStore(cfg *config.Config, store objectstore.ObjectStore, natsConn *nats.Conn, snapshotKV ...jetstream.KeyValue) *VolumeServiceImpl {
 	bucketName := ""
 	if cfg != nil {
 		bucketName = cfg.Predastore.Bucket
@@ -1562,7 +1563,7 @@ func (s *VolumeServiceImpl) DeleteVolume(ctx context.Context, input *ec2.DeleteV
 	// Check if any snapshots reference this volume via JetStream KV.
 	// Snapshot-backed clones read chunk files from the source volume's
 	// S3 prefix via ReadFrom(). Deleting the source would break all clones.
-	if err := s.checkVolumeHasNoSnapshots(volumeID); err != nil {
+	if err := s.checkVolumeHasNoSnapshots(ctx, volumeID); err != nil {
 		return nil, err
 	}
 
@@ -1676,29 +1677,29 @@ func (s *VolumeServiceImpl) getSnapshotMetadata(ctx context.Context, snapshotID
 
 // checkVolumeHasNoSnapshots checks if a volume has dependent snapshots
 // using the JetStream KV index.
-func (s *VolumeServiceImpl) checkVolumeHasNoSnapshots(volumeID string) error {
+func (s *VolumeServiceImpl) checkVolumeHasNoSnapshots(ctx context.Context, volumeID string) error {
 	if s.snapshotKV == nil {
-		slog.Error("checkVolumeHasNoSnapshots: snapshotKV is nil", "volumeId", volumeID)
+		slog.ErrorContext(ctx, "checkVolumeHasNoSnapshots: snapshotKV is nil", "volumeId", volumeID)
 		return errors.New(awserrors.ErrorServerInternal)
 	}
 
-	has, err := s.volumeHasSnapshotsKV(volumeID)
+	has, err := s.volumeHasSnapshotsKV(ctx, volumeID)
 	if err != nil {
-		slog.Error("checkVolumeHasNoSnapshots: KV lookup failed", "volumeId", volumeID, "err", err)
+		slog.ErrorContext(ctx, "checkVolumeHasNoSnapshots: KV lookup failed", "volumeId", volumeID, "err", err)
 		return errors.New(awserrors.ErrorServerInternal)
 	}
 	if has {
-		slog.Error("DeleteVolume blocked: volume has snapshots", "volumeId", volumeID)
+		slog.ErrorContext(ctx, "DeleteVolume blocked: volume has snapshots", "volumeId", volumeID)
 		return errors.New(awserrors.ErrorVolumeInUse)
 	}
 	return nil
 }
 
 // volumeHasSnapshotsKV checks the JetStream KV index for snapshot references.
-func (s *VolumeServiceImpl) volumeHasSnapshotsKV(volumeID string) (bool, error) {
-	entry, err := s.snapshotKV.Get(volumeID)
+func (s *VolumeServiceImpl) volumeHasSnapshotsKV(ctx context.Context, volumeID string) (bool, error) {
+	entry, err := s.snapshotKV.Get(ctx, volumeID)
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return false, nil
 		}
 		return false, err

spinifex/handlers/ec2/volume/service_impl_test.go

@@ -14,10 +14,12 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
+	"github.com/mulgadc/spinifex/spinifex/testutil"
 	"github.com/mulgadc/spinifex/spinifex/types"
 	"github.com/mulgadc/viperblock/viperblock"
 	"github.com/nats-io/nats-server/v2/server"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -510,30 +512,12 @@ func TestCreateVolume_FromSnapshot_CorruptMetadata(t *testing.T) {
 }
 
 // setupTestVolumeKV creates a NATS JetStream test server and returns a KV bucket.
-func setupTestVolumeKV(t *testing.T) nats.KeyValue {
+func setupTestVolumeKV(t *testing.T) jetstream.KeyValue {
 	t.Helper()
-	opts := &server.Options{
-		Host:      "127.0.0.1",
-		Port:      -1,
-		JetStream: true,
-		StoreDir:  t.TempDir(),
-		NoLog:     true,
-		NoSigs:    true,
-	}
-	ns, err := server.NewServer(opts)
-	require.NoError(t, err)
-	go ns.Start()
-	require.True(t, ns.ReadyForConnections(5*time.Second))
-	t.Cleanup(func() { ns.Shutdown() })
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 
-	nc, err := nats.Connect(ns.ClientURL())
-	require.NoError(t, err)
-	t.Cleanup(func() { nc.Close() })
-
-	js, err := nc.JetStream()
-	require.NoError(t, err)
-
-	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{
+	kv, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{
 		Bucket: "spinifex-volume-snapshots",
 	})
 	require.NoError(t, err)
@@ -574,7 +558,7 @@ func TestDeleteVolume_BlockedByKV(t *testing.T) {
 	// Put a snapshot ref in KV
 	data, err := json.Marshal([]string{"snap-001"})
 	require.NoError(t, err)
-	_, err = kv.Put(volumeID, data)
+	_, err = kv.Put(t.Context(), volumeID, data)
 	require.NoError(t, err)
 
 	// DeleteVolume should be blocked
@@ -1479,7 +1463,7 @@ func TestDeleteVolumeOnTerminate_SurfacesDeleteFailure(t *testing.T) {
 	// clear has already succeeded.
 	snapData, err := json.Marshal([]string{"snap-001"})
 	require.NoError(t, err)
-	_, err = kv.Put(volumeID, snapData)
+	_, err = kv.Put(t.Context(), volumeID, snapData)
 	require.NoError(t, err)
 
 	err = svc.DeleteVolumeOnTerminate(context.Background(), volumeID, "")

spinifex/handlers/ec2/vpc/eni.go

@@ -15,7 +15,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
 	"github.com/mulgadc/spinifex/spinifex/network/topology"
 	"github.com/mulgadc/spinifex/spinifex/utils"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 const (
@@ -91,7 +91,7 @@ func (s *VPCServiceImpl) CreateNetworkInterface(ctx context.Context, input *ec2.
 	subnetId := *input.SubnetId
 
 	// Verify subnet exists and belongs to this account
-	subnetEntry, err := s.subnetKV.Get(utils.AccountKey(accountID, subnetId))
+	subnetEntry, err := s.subnetKV.Get(ctx, utils.AccountKey(accountID, subnetId))
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidSubnetIDNotFound)
 	}
@@ -109,13 +109,13 @@ func (s *VPCServiceImpl) CreateNetworkInterface(ctx context.Context, input *ec2.
 	}
 	if len(sgIdsIn) == 0 {
 		// AWS attaches the per-VPC default SG when the caller omits Groups.
-		defaultSGId, err := s.FindDefaultSGForVPC(accountID, subnet.VpcId)
+		defaultSGId, err := s.findDefaultSGForVPC(ctx, accountID, subnet.VpcId)
 		if err != nil || defaultSGId == "" {
 			return nil, errors.New(awserrors.ErrorServerInternal)
 		}
 		sgIdsIn = []string{defaultSGId}
 	}
-	if err := s.validateSGAttachment(accountID, sgIdsIn, subnet.VpcId); err != nil {
+	if err := s.validateSGAttachment(ctx, accountID, sgIdsIn, subnet.VpcId); err != nil {
 		return nil, err
 	}
 
@@ -127,7 +127,7 @@ func (s *VPCServiceImpl) CreateNetworkInterface(ctx context.Context, input *ec2.
 		// TODO: validate the requested IP is in the subnet range and not already allocated
 		privateIP = *input.PrivateIpAddress
 	} else {
-		ip, err := s.ipam.AllocateIP(subnetId, subnet.CidrBlock, PurposeENIPrimary, eniId)
+		ip, err := s.ipam.AllocateIP(ctx, subnetId, subnet.CidrBlock, PurposeENIPrimary, eniId)
 		if err != nil {
 			return nil, errors.New(awserrors.ErrorServerInternal)
 		}
@@ -160,7 +160,7 @@ func (s *VPCServiceImpl) CreateNetworkInterface(ctx context.Context, input *ec2.
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal ENI record: %w", err)
 	}
-	if _, err := s.eniKV.Put(utils.AccountKey(accountID, eniId), data); err != nil {
+	if _, err := s.eniKV.Put(ctx, utils.AccountKey(accountID, eniId), data); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -205,9 +205,9 @@ func (s *VPCServiceImpl) deleteNetworkInterface(ctx context.Context, eniId, acco
 	key := utils.AccountKey(accountID, eniId)
 
 	// Get the ENI record
-	entry, err := s.eniKV.Get(key)
+	entry, err := s.eniKV.Get(ctx, key)
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			// Internal teardown (force) tolerates an already-gone ENI so
 			// instance terminate converges (ADR-0003 §2); the public API is
 			// AWS-faithful and returns NotFound.
@@ -232,7 +232,7 @@ func (s *VPCServiceImpl) deleteNetworkInterface(ctx context.Context, eniId, acco
 	}
 
 	// Release the private IP back to the IPAM pool
-	if err := s.ipam.ReleaseIP(record.SubnetId, record.PrivateIpAddress); err != nil {
+	if err := s.ipam.ReleaseIP(ctx, record.SubnetId, record.PrivateIpAddress); err != nil {
 		slog.WarnContext(ctx, "Failed to release IP during ENI delete", "eni", eniId, "ip", record.PrivateIpAddress, "err", err)
 	}
 
@@ -255,7 +255,7 @@ func (s *VPCServiceImpl) deleteNetworkInterface(ctx context.Context, eniId, acco
 		}
 	}
 
-	if err := s.eniKV.Delete(key); err != nil {
+	if err := s.eniKV.Delete(ctx, key); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -286,7 +286,7 @@ func (s *VPCServiceImpl) ModifyNetworkInterfaceAttribute(ctx context.Context, in
 	eniId := *input.NetworkInterfaceId
 	key := utils.AccountKey(accountID, eniId)
 
-	entry, err := s.eniKV.Get(key)
+	entry, err := s.eniKV.Get(ctx, key)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidNetworkInterfaceIDNotFound)
 	}
@@ -305,7 +305,7 @@ func (s *VPCServiceImpl) ModifyNetworkInterfaceAttribute(ctx context.Context, in
 				sgIds = append(sgIds, *id)
 			}
 		}
-		if err := s.validateSGAttachment(accountID, sgIds, record.VpcId); err != nil {
+		if err := s.validateSGAttachment(ctx, accountID, sgIds, record.VpcId); err != nil {
 			return nil, err
 		}
 		record.SecurityGroupIds = sgIds
@@ -321,7 +321,7 @@ func (s *VPCServiceImpl) ModifyNetworkInterfaceAttribute(ctx context.Context, in
 		slog.ErrorContext(ctx, "ModifyNetworkInterfaceAttribute: failed to marshal ENI record", "eniId", eniId, "accountID", accountID, "err", err)
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.eniKV.Update(key, data, entry.Revision()); err != nil {
+	if _, err := s.eniKV.Update(ctx, key, data, entry.Revision()); err != nil {
 		slog.ErrorContext(ctx, "ModifyNetworkInterfaceAttribute: KV update failed", "eniId", eniId, "accountID", accountID, "err", err)
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
@@ -371,8 +371,8 @@ func (s *VPCServiceImpl) DescribeNetworkInterfaces(ctx context.Context, input *e
 	}
 
 	prefix := accountID + "."
-	keys, err := s.eniKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	keys, err := s.eniKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -384,7 +384,7 @@ func (s *VPCServiceImpl) DescribeNetworkInterfaces(ctx context.Context, input *e
 			continue
 		}
 
-		entry, err := s.eniKV.Get(key)
+		entry, err := s.eniKV.Get(ctx, key)
 		if err != nil {
 			slog.WarnContext(ctx, "Failed to get ENI record", "key", key, "error", err)
 			continue
@@ -504,8 +504,12 @@ func eniMatchesFilters(record *ENIRecord, filters map[string][]string) bool {
 // AttachENI marks an ENI as attached to an instance (internal use by RunInstances).
 // accountID scopes the lookup to the correct KV key.
 func (s *VPCServiceImpl) AttachENI(accountID, eniId, instanceId string, deviceIndex int64) (string, error) {
+	return s.attachENI(context.Background(), accountID, eniId, instanceId, deviceIndex)
+}
+
+func (s *VPCServiceImpl) attachENI(ctx context.Context, accountID, eniId, instanceId string, deviceIndex int64) (string, error) {
 	key := utils.AccountKey(accountID, eniId)
-	entry, err := s.eniKV.Get(key)
+	entry, err := s.eniKV.Get(ctx, key)
 	if err != nil {
 		return "", errors.New(awserrors.ErrorInvalidNetworkInterfaceIDNotFound)
 	}
@@ -529,11 +533,11 @@ func (s *VPCServiceImpl) AttachENI(accountID, eniId, instanceId string, deviceIn
 	if err != nil {
 		return "", fmt.Errorf("failed to marshal ENI record: %w", err)
 	}
-	if _, err := s.eniKV.Update(key, data, entry.Revision()); err != nil {
+	if _, err := s.eniKV.Update(ctx, key, data, entry.Revision()); err != nil {
 		return "", errors.New(awserrors.ErrorServerInternal)
 	}
 
-	slog.Info("ENI attached", "eniId", eniId, "instanceId", instanceId, "attachmentId", attachmentId)
+	slog.InfoContext(ctx, "ENI attached", "eniId", eniId, "instanceId", instanceId, "attachmentId", attachmentId)
 	return attachmentId, nil
 }
 
@@ -541,7 +545,7 @@ func (s *VPCServiceImpl) AttachENI(accountID, eniId, instanceId string, deviceIn
 // accountID scopes the lookup to the correct KV key.
 func (s *VPCServiceImpl) DetachENI(ctx context.Context, accountID, eniId string) error {
 	key := utils.AccountKey(accountID, eniId)
-	entry, err := s.eniKV.Get(key)
+	entry, err := s.eniKV.Get(ctx, key)
 	if err != nil {
 		return errors.New(awserrors.ErrorInvalidNetworkInterfaceIDNotFound)
 	}
@@ -560,7 +564,7 @@ func (s *VPCServiceImpl) DetachENI(ctx context.Context, accountID, eniId string)
 	if err != nil {
 		return fmt.Errorf("failed to marshal ENI record: %w", err)
 	}
-	if _, err := s.eniKV.Update(key, data, entry.Revision()); err != nil {
+	if _, err := s.eniKV.Update(ctx, key, data, entry.Revision()); err != nil {
 		return errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -572,8 +576,12 @@ func (s *VPCServiceImpl) DetachENI(ctx context.Context, accountID, eniId string)
 // hot-plug handlers that need the record's MAC + AttachmentId before
 // running the QMP pipeline.
 func (s *VPCServiceImpl) GetENIRecord(accountID, eniId string) (ENIRecord, error) {
+	return s.getENIRecord(context.Background(), accountID, eniId)
+}
+
+func (s *VPCServiceImpl) getENIRecord(ctx context.Context, accountID, eniId string) (ENIRecord, error) {
 	key := utils.AccountKey(accountID, eniId)
-	entry, err := s.eniKV.Get(key)
+	entry, err := s.eniKV.Get(ctx, key)
 	if err != nil {
 		return ENIRecord{}, errors.New(awserrors.ErrorInvalidNetworkInterfaceIDNotFound)
 	}
@@ -588,8 +596,12 @@ func (s *VPCServiceImpl) GetENIRecord(accountID, eniId string) (ENIRecord, error
 // the result under the same revision token (compare-and-swap). fn must
 // not block on external I/O; the KV update slot is contended.
 func (s *VPCServiceImpl) UpdateENI(accountID, eniId string, fn func(*ENIRecord)) error {
+	return s.updateENI(context.Background(), accountID, eniId, fn)
+}
+
+func (s *VPCServiceImpl) updateENI(ctx context.Context, accountID, eniId string, fn func(*ENIRecord)) error {
 	key := utils.AccountKey(accountID, eniId)
-	entry, err := s.eniKV.Get(key)
+	entry, err := s.eniKV.Get(ctx, key)
 	if err != nil {
 		return errors.New(awserrors.ErrorInvalidNetworkInterfaceIDNotFound)
 	}
@@ -602,7 +614,7 @@ func (s *VPCServiceImpl) UpdateENI(accountID, eniId string, fn func(*ENIRecord))
 	if err != nil {
 		return fmt.Errorf("marshal ENI record: %w", err)
 	}
-	if _, err := s.eniKV.Update(key, data, entry.Revision()); err != nil {
+	if _, err := s.eniKV.Update(ctx, key, data, entry.Revision()); err != nil {
 		return errors.New(awserrors.ErrorServerInternal)
 	}
 	return nil
@@ -612,9 +624,13 @@ func (s *VPCServiceImpl) UpdateENI(accountID, eniId string, fn func(*ENIRecord))
 // instanceID. Used by the hot-plug reconciler to converge a node's instances
 // against KV on restart.
 func (s *VPCServiceImpl) ListInstanceENIs(accountID, instanceID string) ([]ENIRecord, error) {
-	keys, err := s.eniKV.Keys()
+	return s.listInstanceENIs(context.Background(), accountID, instanceID)
+}
+
+func (s *VPCServiceImpl) listInstanceENIs(ctx context.Context, accountID, instanceID string) ([]ENIRecord, error) {
+	keys, err := s.eniKV.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, nil
 		}
 		return nil, errors.New(awserrors.ErrorServerInternal)
@@ -625,7 +641,7 @@ func (s *VPCServiceImpl) ListInstanceENIs(accountID, instanceID string) ([]ENIRe
 		if key == utils.VersionKey || !strings.HasPrefix(key, prefix) {
 			continue
 		}
-		entry, err := s.eniKV.Get(key)
+		entry, err := s.eniKV.Get(ctx, key)
 		if err != nil {
 			continue
 		}
@@ -643,8 +659,12 @@ func (s *VPCServiceImpl) ListInstanceENIs(accountID, instanceID string) ([]ENIRe
 // FindENIByAttachment scans the ENI bucket for the record with the given
 // AttachmentId. Used by DetachNetworkInterface which identifies by attachment ID.
 func (s *VPCServiceImpl) FindENIByAttachment(accountID, attachmentId string) (ENIRecord, error) {
-	keys, err := s.eniKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	return s.findENIByAttachment(context.Background(), accountID, attachmentId)
+}
+
+func (s *VPCServiceImpl) findENIByAttachment(ctx context.Context, accountID, attachmentId string) (ENIRecord, error) {
+	keys, err := s.eniKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return ENIRecord{}, errors.New(awserrors.ErrorServerInternal)
 	}
 	prefix := accountID + "."
@@ -652,7 +672,7 @@ func (s *VPCServiceImpl) FindENIByAttachment(accountID, attachmentId string) (EN
 		if key == utils.VersionKey || !strings.HasPrefix(key, prefix) {
 			continue
 		}
-		entry, err := s.eniKV.Get(key)
+		entry, err := s.eniKV.Get(ctx, key)
 		if err != nil {
 			continue
 		}
@@ -669,8 +689,12 @@ func (s *VPCServiceImpl) FindENIByAttachment(accountID, attachmentId string) (EN
 
 // UpdateENIPublicIP updates the PublicIpAddress and PublicIpPool on an ENI record.
 func (s *VPCServiceImpl) UpdateENIPublicIP(accountID, eniId, publicIP, poolName string) error {
+	return s.updateENIPublicIP(context.Background(), accountID, eniId, publicIP, poolName)
+}
+
+func (s *VPCServiceImpl) updateENIPublicIP(ctx context.Context, accountID, eniId, publicIP, poolName string) error {
 	key := utils.AccountKey(accountID, eniId)
-	entry, err := s.eniKV.Get(key)
+	entry, err := s.eniKV.Get(ctx, key)
 	if err != nil {
 		return fmt.Errorf("ENI %s not found: %w", eniId, err)
 	}
@@ -687,11 +711,11 @@ func (s *VPCServiceImpl) UpdateENIPublicIP(accountID, eniId, publicIP, poolName
 	if err != nil {
 		return fmt.Errorf("marshal ENI record: %w", err)
 	}
-	if _, err := s.eniKV.Update(key, data, entry.Revision()); err != nil {
+	if _, err := s.eniKV.Update(ctx, key, data, entry.Revision()); err != nil {
 		return fmt.Errorf("update ENI record: %w", err)
 	}
 
-	slog.Info("Updated ENI with public IP", "eniId", eniId, "publicIp", publicIP, "pool", poolName)
+	slog.InfoContext(ctx, "Updated ENI with public IP", "eniId", eniId, "publicIp", publicIP, "pool", poolName)
 	return nil
 }
 
@@ -832,7 +856,7 @@ func (s *VPCServiceImpl) publishNATEvent(topic, vpcId, externalIP, logicalIP, po
 // or MissingParameter per AWS contract.
 const sgPerInterfaceLimit = 5
 
-func (s *VPCServiceImpl) validateSGAttachment(accountID string, sgIds []string, vpcId string) error {
+func (s *VPCServiceImpl) validateSGAttachment(ctx context.Context, accountID string, sgIds []string, vpcId string) error {
 	if len(sgIds) == 0 {
 		return errors.New(awserrors.ErrorMissingParameter)
 	}
@@ -840,13 +864,13 @@ func (s *VPCServiceImpl) validateSGAttachment(accountID string, sgIds []string,
 		return errors.New(awserrors.ErrorSecurityGroupsPerInterfaceLimitExceeded)
 	}
 
-	if err := s.requireVPCExists(accountID, vpcId); err != nil {
+	if err := s.requireVPCExists(ctx, accountID, vpcId); err != nil {
 		return err
 	}
 
 	// Each SG must exist in the caller's account and belong to the same VPC.
 	for _, sgId := range sgIds {
-		sgEntry, err := s.sgKV.Get(utils.AccountKey(accountID, sgId))
+		sgEntry, err := s.sgKV.Get(ctx, utils.AccountKey(accountID, sgId))
 		if err != nil {
 			return errors.New(awserrors.ErrorInvalidGroupNotFound)
 		}
@@ -868,9 +892,9 @@ func (s *VPCServiceImpl) isEIPOwned(ctx context.Context, eniId, accountID string
 	if s.eipKV == nil {
 		return false, nil
 	}
-	keys, err := s.eipKV.Keys()
+	keys, err := s.eipKV.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return false, nil
 		}
 		return false, fmt.Errorf("eipKV.Keys: %w", err)
@@ -880,7 +904,7 @@ func (s *VPCServiceImpl) isEIPOwned(ctx context.Context, eniId, accountID string
 		if len(k) < len(prefix) || k[:len(prefix)] != prefix {
 			continue
 		}
-		entry, err := s.eipKV.Get(k)
+		entry, err := s.eipKV.Get(ctx, k)
 		if err != nil {
 			return false, fmt.Errorf("eipKV.Get(%s): %w", k, err)
 		}

spinifex/handlers/ec2/vpc/eni_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/testutil"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -935,27 +936,28 @@ func TestDescribeNetworkInterfaces_FilterByAttachmentStatus(t *testing.T) {
 // --- isEIPOwned tests ---
 
 // setupEIPTestService creates a VPC service with an EIP KV bucket wired in.
-func setupEIPTestService(t *testing.T) (*VPCServiceImpl, nats.KeyValue) {
+func setupEIPTestService(t *testing.T) (*VPCServiceImpl, jetstream.KeyValue) {
 	t.Helper()
-	_, nc, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 
-	svc, err := NewVPCServiceImplWithNATS(nil, nc)
+	svc, err := NewVPCServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 
-	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{Bucket: "eip-test"})
+	kv, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{Bucket: "eip-test"})
 	require.NoError(t, err)
 	svc.eipKV = kv
 
 	return svc, kv
 }
 
-func putEIPRecord(t *testing.T, kv nats.KeyValue, key, eniID string) {
+func putEIPRecord(t *testing.T, kv jetstream.KeyValue, key, eniID string) {
 	t.Helper()
 	data, err := json.Marshal(struct {
 		ENIId string `json:"eni_id"`
 	}{ENIId: eniID})
 	require.NoError(t, err)
-	_, err = kv.Put(key, data)
+	_, err = kv.Put(t.Context(), key, data)
 	require.NoError(t, err)
 }
 
@@ -1004,7 +1006,7 @@ func TestIsEIPOwned_WrongAccountPrefix(t *testing.T) {
 
 func TestIsEIPOwned_MalformedJSON(t *testing.T) {
 	svc, kv := setupEIPTestService(t)
-	_, err := kv.Put(testAccountID+".eipalloc-001", []byte("not json"))
+	_, err := kv.Put(t.Context(), testAccountID+".eipalloc-001", []byte("not json"))
 	require.NoError(t, err)
 
 	owned, err := svc.isEIPOwned(context.Background(), "eni-111", testAccountID)
@@ -1030,7 +1032,7 @@ func TestValidateSGAttachment_OK(t *testing.T) {
 	vpcID := createTestVPC(t, svc, "10.0.0.0/16")
 	sg := createTestSG(t, svc, vpcID, "ok")
 
-	err := svc.validateSGAttachment(testAccountID, []string{sg}, vpcID)
+	err := svc.validateSGAttachment(t.Context(), testAccountID, []string{sg}, vpcID)
 	assert.NoError(t, err)
 }
 
@@ -1038,7 +1040,7 @@ func TestValidateSGAttachment_EmptyList(t *testing.T) {
 	svc := setupTestVPCService(t)
 	vpcID := createTestVPC(t, svc, "10.0.0.0/16")
 
-	err := svc.validateSGAttachment(testAccountID, nil, vpcID)
+	err := svc.validateSGAttachment(t.Context(), testAccountID, nil, vpcID)
 	assert.ErrorContains(t, err, "MissingParameter")
 }
 
@@ -1050,13 +1052,13 @@ func TestValidateSGAttachment_TooMany(t *testing.T) {
 		sgs[i] = createTestSG(t, svc, vpcID, fmt.Sprintf("sg-%d", i))
 	}
 
-	err := svc.validateSGAttachment(testAccountID, sgs, vpcID)
+	err := svc.validateSGAttachment(t.Context(), testAccountID, sgs, vpcID)
 	assert.ErrorContains(t, err, "SecurityGroupsPerInterfaceLimitExceeded")
 }
 
 func TestValidateSGAttachment_UnknownVPC(t *testing.T) {
 	svc := setupTestVPCService(t)
-	err := svc.validateSGAttachment(testAccountID, []string{"sg-aaa"}, "vpc-missing")
+	err := svc.validateSGAttachment(t.Context(), testAccountID, []string{"sg-aaa"}, "vpc-missing")
 	assert.ErrorContains(t, err, "InvalidVpcID.NotFound")
 }
 
@@ -1064,7 +1066,7 @@ func TestValidateSGAttachment_SGNotFound(t *testing.T) {
 	svc := setupTestVPCService(t)
 	vpcID := createTestVPC(t, svc, "10.0.0.0/16")
 
-	err := svc.validateSGAttachment(testAccountID, []string{"sg-doesntexist"}, vpcID)
+	err := svc.validateSGAttachment(t.Context(), testAccountID, []string{"sg-doesntexist"}, vpcID)
 	assert.ErrorContains(t, err, "InvalidGroup.NotFound")
 }
 
@@ -1074,7 +1076,7 @@ func TestValidateSGAttachment_CrossVPC(t *testing.T) {
 	vpcB := createTestVPC(t, svc, "10.1.0.0/16")
 	sgInA := createTestSG(t, svc, vpcA, "in-a")
 
-	err := svc.validateSGAttachment(testAccountID, []string{sgInA}, vpcB)
+	err := svc.validateSGAttachment(t.Context(), testAccountID, []string{sgInA}, vpcB)
 	assert.ErrorContains(t, err, "InvalidParameterValue")
 }
 

spinifex/handlers/ec2/vpc/external_ipam.go

@@ -10,7 +10,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/network/external"
 	"github.com/mulgadc/spinifex/spinifex/network/external/dhcp"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Backward-compatible aliases so callers (daemon, EIP, EC2 handlers, tests)
@@ -29,7 +29,7 @@ const (
 // ExternalIPAM is the AWS-facing entry point for external IP allocation,
 // dispatching to StaticPoolAllocator or dhcp.DHCPPoolAllocator per pool name.
 type ExternalIPAM struct {
-	kv      nats.KeyValue
+	kv      jetstream.KeyValue
 	pools   []external.ExternalPoolConfig
 	static  *external.StaticPoolAllocator
 	perPool map[string]external.Allocator // dhcp overrides; static pools fall through to `static`
@@ -38,15 +38,15 @@ type ExternalIPAM struct {
 // NewExternalIPAM creates a new ExternalIPAM. Static pools wire through
 // external.StaticPoolAllocator; DHCP-sourced pools wait for EnableDHCP
 // to install the per-pool dhcp.DHCPPoolAllocator.
-func NewExternalIPAM(js nats.JetStreamContext, pools []external.ExternalPoolConfig) (*ExternalIPAM, error) {
+func NewExternalIPAM(ctx context.Context, js jetstream.JetStream, pools []external.ExternalPoolConfig) (*ExternalIPAM, error) {
 	staticPools := filterStatic(pools)
 	var (
 		alloc *external.StaticPoolAllocator
-		kv    nats.KeyValue
+		kv    jetstream.KeyValue
 	)
 	if len(staticPools) > 0 {
 		var err error
-		alloc, err = external.NewStaticPoolAllocator(js, staticPools)
+		alloc, err = external.NewStaticPoolAllocator(ctx, js, staticPools)
 		if err != nil {
 			return nil, err
 		}
@@ -56,7 +56,7 @@ func NewExternalIPAM(js nats.JetStreamContext, pools []external.ExternalPoolConf
 }
 
 // NewExternalIPAMWithKV creates an ExternalIPAM with an existing KV bucket (for testing).
-func NewExternalIPAMWithKV(kv nats.KeyValue, pools []external.ExternalPoolConfig) *ExternalIPAM {
+func NewExternalIPAMWithKV(kv jetstream.KeyValue, pools []external.ExternalPoolConfig) *ExternalIPAM {
 	alloc := external.NewStaticPoolAllocatorWithKV(kv, filterStatic(pools))
 	return &ExternalIPAM{kv: kv, pools: pools, static: alloc, perPool: map[string]external.Allocator{}}
 }
@@ -128,10 +128,14 @@ func (m *ExternalIPAM) ReleaseIP(ctx context.Context, poolName, ip, ownerENIID s
 // GetPoolRecord returns the current IPAM record for a pool. DHCP-sourced
 // pools have no static record — the per-AZ lease bucket is authoritative.
 func (m *ExternalIPAM) GetPoolRecord(poolName string) (*ExternalIPAMRecord, error) {
+	return m.getPoolRecord(context.Background(), poolName)
+}
+
+func (m *ExternalIPAM) getPoolRecord(ctx context.Context, poolName string) (*ExternalIPAMRecord, error) {
 	if m.static == nil {
 		return nil, fmt.Errorf("pool record unavailable: no static allocator")
 	}
-	return m.static.GetPoolRecord(poolName)
+	return m.static.GetPoolRecord(ctx, poolName)
 }
 
 func (m *ExternalIPAM) allocatorFor(poolName string) (external.Allocator, error) {

spinifex/handlers/ec2/vpc/external_ipam_test.go

@@ -14,9 +14,10 @@ import (
 
 func setupTestExternalIPAM(t *testing.T, pools []external.ExternalPoolConfig) *ExternalIPAM {
 	t.Helper()
-	_, _, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 
-	ipam, err := NewExternalIPAM(js, pools)
+	ipam, err := NewExternalIPAM(t.Context(), js, pools)
 	require.NoError(t, err)
 	return ipam
 }
@@ -338,10 +339,11 @@ func TestExternalIPAM_RangeValidation(t *testing.T) {
 func TestExternalIPAM_InitFromConfig(t *testing.T) {
 	pool := testPool()
 	// Create IPAM twice — second init should be idempotent
-	_, _, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 
 	// First init
-	ipam1, err := NewExternalIPAM(js, []external.ExternalPoolConfig{pool})
+	ipam1, err := NewExternalIPAM(t.Context(), js, []external.ExternalPoolConfig{pool})
 	require.NoError(t, err)
 
 	// Allocate an IP
@@ -350,7 +352,7 @@ func TestExternalIPAM_InitFromConfig(t *testing.T) {
 	assert.Equal(t, "192.168.1.151", ip)
 
 	// Second init (simulating restart) — should not lose allocation
-	ipam2, err := NewExternalIPAM(js, []external.ExternalPoolConfig{pool})
+	ipam2, err := NewExternalIPAM(t.Context(), js, []external.ExternalPoolConfig{pool})
 	require.NoError(t, err)
 
 	record, err := ipam2.GetPoolRecord("wan")

spinifex/handlers/ec2/vpc/ipam.go

@@ -1,15 +1,16 @@
 package handlers_ec2_vpc
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"log/slog"
 	"net/netip"
 
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
-	"github.com/mulgadc/spinifex/spinifex/utils"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 const (
@@ -34,32 +35,32 @@ type IPAMRecord struct {
 
 // IPAM manages IP address allocation for VPC subnets using NATS KV with CAS.
 type IPAM struct {
-	kv nats.KeyValue
+	kv jetstream.KeyValue
 }
 
 // NewIPAM creates a new IPAM instance backed by NATS JetStream KV.
-func NewIPAM(js nats.JetStreamContext) (*IPAM, error) {
-	kv, err := utils.GetOrCreateKVBucket(js, KVBucketIPAM, 5)
+func NewIPAM(ctx context.Context, js jetstream.JetStream) (*IPAM, error) {
+	kv, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketIPAM, 5)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create IPAM KV bucket: %w", err)
 	}
-	if err := migrate.DefaultRegistry.RunKVWithJetStream(KVBucketIPAM, kv, js, KVBucketIPAMVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKVWithJetStream(ctx, KVBucketIPAM, kv, js, KVBucketIPAMVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketIPAM, err)
 	}
 	return &IPAM{kv: kv}, nil
 }
 
 // NewIPAMWithKV creates an IPAM with an existing KV bucket (for testing).
-func NewIPAMWithKV(kv nats.KeyValue) *IPAM {
+func NewIPAMWithKV(kv jetstream.KeyValue) *IPAM {
 	return &IPAM{kv: kv}
 }
 
 // AllocateIP allocates an IP from the subnet, reserving the first 4 and last
 // addresses per AWS convention. Uses CAS for conflict-free multi-node allocation.
-func (m *IPAM) AllocateIP(subnetId, cidrBlock, purpose, ownerID string) (string, error) {
+func (m *IPAM) AllocateIP(ctx context.Context, subnetId, cidrBlock, purpose, ownerID string) (string, error) {
 	for attempt := range 5 {
-		record, revision, err := m.getRecord(subnetId)
-		if err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
+		record, revision, err := m.getRecord(ctx, subnetId)
+		if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
 			return "", fmt.Errorf("get IPAM record: %w", err)
 		}
 
@@ -85,9 +86,9 @@ func (m *IPAM) AllocateIP(subnetId, cidrBlock, purpose, ownerID string) (string,
 		}
 
 		if revision == 0 {
-			_, err = m.kv.Create(subnetId, data)
+			_, err = m.kv.Create(ctx, subnetId, data)
 		} else {
-			_, err = m.kv.Update(subnetId, data, revision)
+			_, err = m.kv.Update(ctx, subnetId, data, revision)
 		}
 
 		if err != nil {
@@ -103,9 +104,9 @@ func (m *IPAM) AllocateIP(subnetId, cidrBlock, purpose, ownerID string) (string,
 }
 
 // ReleaseIP releases a previously allocated IP address back to the subnet pool.
-func (m *IPAM) ReleaseIP(subnetId, ip string) error {
+func (m *IPAM) ReleaseIP(ctx context.Context, subnetId, ip string) error {
 	for attempt := range 5 {
-		record, revision, err := m.getRecord(subnetId)
+		record, revision, err := m.getRecord(ctx, subnetId)
 		if err != nil {
 			return fmt.Errorf("get IPAM record for release: %w", err)
 		}
@@ -130,7 +131,7 @@ func (m *IPAM) ReleaseIP(subnetId, ip string) error {
 			return fmt.Errorf("marshal IPAM record: %w", err)
 		}
 
-		if _, err := m.kv.Update(subnetId, data, revision); err != nil {
+		if _, err := m.kv.Update(ctx, subnetId, data, revision); err != nil {
 			slog.Debug("IPAM release CAS conflict, retrying", "subnet", subnetId, "attempt", attempt)
 			continue
 		}
@@ -143,10 +144,10 @@ func (m *IPAM) ReleaseIP(subnetId, ip string) error {
 }
 
 // AllocatedIPs returns the list of allocated IP entries for a subnet.
-func (m *IPAM) AllocatedIPs(subnetId string) ([]IPEntry, error) {
-	record, _, err := m.getRecord(subnetId)
+func (m *IPAM) AllocatedIPs(ctx context.Context, subnetId string) ([]IPEntry, error) {
+	record, _, err := m.getRecord(ctx, subnetId)
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return nil, nil
 		}
 		return nil, err
@@ -157,11 +158,11 @@ func (m *IPAM) AllocatedIPs(subnetId string) ([]IPEntry, error) {
 	return record.Allocated, nil
 }
 
-func (m *IPAM) getRecord(subnetId string) (*IPAMRecord, uint64, error) {
-	entry, err := m.kv.Get(subnetId)
+func (m *IPAM) getRecord(ctx context.Context, subnetId string) (*IPAMRecord, uint64, error) {
+	entry, err := m.kv.Get(ctx, subnetId)
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
-			return nil, 0, nats.ErrKeyNotFound
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
+			return nil, 0, jetstream.ErrKeyNotFound
 		}
 		return nil, 0, err
 	}

spinifex/handlers/ec2/vpc/ipam_test.go

@@ -6,16 +6,17 @@ import (
 	"testing"
 
 	"github.com/mulgadc/spinifex/spinifex/testutil"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func setupTestIPAM(t *testing.T) *IPAM {
 	t.Helper()
-	_, _, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 
-	ipam, err := NewIPAM(js)
+	ipam, err := NewIPAM(t.Context(), js)
 	require.NoError(t, err)
 	return ipam
 }
@@ -23,7 +24,7 @@ func setupTestIPAM(t *testing.T) *IPAM {
 func TestIPAM_AllocateFirst(t *testing.T) {
 	ipam := setupTestIPAM(t)
 
-	ip, err := ipam.AllocateIP("subnet-1", "10.0.1.0/24", PurposeENIPrimary, "eni-1")
+	ip, err := ipam.AllocateIP(t.Context(), "subnet-1", "10.0.1.0/24", PurposeENIPrimary, "eni-1")
 	require.NoError(t, err)
 	// First allocable IP is .4 (skip .0 network, .1 gateway, .2 DNS, .3 reserved)
 	assert.Equal(t, "10.0.1.4", ip)
@@ -32,15 +33,15 @@ func TestIPAM_AllocateFirst(t *testing.T) {
 func TestIPAM_AllocateSequential(t *testing.T) {
 	ipam := setupTestIPAM(t)
 
-	ip1, err := ipam.AllocateIP("subnet-seq", "10.0.2.0/24", PurposeENIPrimary, "eni-seq")
+	ip1, err := ipam.AllocateIP(t.Context(), "subnet-seq", "10.0.2.0/24", PurposeENIPrimary, "eni-seq")
 	require.NoError(t, err)
 	assert.Equal(t, "10.0.2.4", ip1)
 
-	ip2, err := ipam.AllocateIP("subnet-seq", "10.0.2.0/24", PurposeENIPrimary, "eni-seq")
+	ip2, err := ipam.AllocateIP(t.Context(), "subnet-seq", "10.0.2.0/24", PurposeENIPrimary, "eni-seq")
 	require.NoError(t, err)
 	assert.Equal(t, "10.0.2.5", ip2)
 
-	ip3, err := ipam.AllocateIP("subnet-seq", "10.0.2.0/24", PurposeENIPrimary, "eni-seq")
+	ip3, err := ipam.AllocateIP(t.Context(), "subnet-seq", "10.0.2.0/24", PurposeENIPrimary, "eni-seq")
 	require.NoError(t, err)
 	assert.Equal(t, "10.0.2.6", ip3)
 }
@@ -48,20 +49,20 @@ func TestIPAM_AllocateSequential(t *testing.T) {
 func TestIPAM_Release(t *testing.T) {
 	ipam := setupTestIPAM(t)
 
-	ip1, err := ipam.AllocateIP("subnet-rel", "10.0.3.0/24", PurposeENIPrimary, "eni-rel")
+	ip1, err := ipam.AllocateIP(t.Context(), "subnet-rel", "10.0.3.0/24", PurposeENIPrimary, "eni-rel")
 	require.NoError(t, err)
 	assert.Equal(t, "10.0.3.4", ip1)
 
-	ip2, err := ipam.AllocateIP("subnet-rel", "10.0.3.0/24", PurposeENIPrimary, "eni-rel")
+	ip2, err := ipam.AllocateIP(t.Context(), "subnet-rel", "10.0.3.0/24", PurposeENIPrimary, "eni-rel")
 	require.NoError(t, err)
 	assert.Equal(t, "10.0.3.5", ip2)
 
 	// Release first IP
-	err = ipam.ReleaseIP("subnet-rel", "10.0.3.4")
+	err = ipam.ReleaseIP(t.Context(), "subnet-rel", "10.0.3.4")
 	require.NoError(t, err)
 
 	// Next allocation should reuse the released IP
-	ip3, err := ipam.AllocateIP("subnet-rel", "10.0.3.0/24", PurposeENIPrimary, "eni-rel")
+	ip3, err := ipam.AllocateIP(t.Context(), "subnet-rel", "10.0.3.0/24", PurposeENIPrimary, "eni-rel")
 	require.NoError(t, err)
 	assert.Equal(t, "10.0.3.4", ip3)
 }
@@ -69,17 +70,17 @@ func TestIPAM_Release(t *testing.T) {
 func TestIPAM_ReleaseNotAllocated(t *testing.T) {
 	ipam := setupTestIPAM(t)
 
-	_, err := ipam.AllocateIP("subnet-rna", "10.0.4.0/24", PurposeENIPrimary, "eni-rna")
+	_, err := ipam.AllocateIP(t.Context(), "subnet-rna", "10.0.4.0/24", PurposeENIPrimary, "eni-rna")
 	require.NoError(t, err)
 
-	err = ipam.ReleaseIP("subnet-rna", "10.0.4.99")
+	err = ipam.ReleaseIP(t.Context(), "subnet-rna", "10.0.4.99")
 	assert.ErrorContains(t, err, "not allocated")
 }
 
 func TestIPAM_ReleaseNoRecord(t *testing.T) {
 	ipam := setupTestIPAM(t)
 
-	err := ipam.ReleaseIP("subnet-nonexistent", "10.0.0.4")
+	err := ipam.ReleaseIP(t.Context(), "subnet-nonexistent", "10.0.0.4")
 	assert.Error(t, err)
 }
 
@@ -92,7 +93,7 @@ func TestIPAM_Exhaustion(t *testing.T) {
 
 	var allocated []string
 	for i := range 11 {
-		ip, err := ipam.AllocateIP(subnetId, cidr, PurposeENIPrimary, "eni-test")
+		ip, err := ipam.AllocateIP(t.Context(), subnetId, cidr, PurposeENIPrimary, "eni-test")
 		require.NoError(t, err, "allocation %d should succeed", i)
 		allocated = append(allocated, ip)
 	}
@@ -102,7 +103,7 @@ func TestIPAM_Exhaustion(t *testing.T) {
 	assert.Equal(t, "10.0.5.14", allocated[len(allocated)-1])
 
 	// Next allocation should fail — subnet exhausted
-	_, err := ipam.AllocateIP(subnetId, cidr, PurposeENIPrimary, "eni-test")
+	_, err := ipam.AllocateIP(t.Context(), subnetId, cidr, PurposeENIPrimary, "eni-test")
 	assert.ErrorContains(t, err, "exhausted")
 }
 
@@ -114,7 +115,7 @@ func TestIPAM_ReservedIPs(t *testing.T) {
 
 	// First 4 allocations should be .4, .5, .6, .7 (skipping .0-.3)
 	for i := 4; i <= 7; i++ {
-		ip, err := ipam.AllocateIP(subnetId, cidr, PurposeENIPrimary, "eni-test")
+		ip, err := ipam.AllocateIP(t.Context(), subnetId, cidr, PurposeENIPrimary, "eni-test")
 		require.NoError(t, err)
 		expected := "10.0.6." + itoa(i)
 		assert.Equal(t, expected, ip)
@@ -125,15 +126,15 @@ func TestIPAM_AllocatedIPs(t *testing.T) {
 	ipam := setupTestIPAM(t)
 
 	// No allocations yet
-	ips, err := ipam.AllocatedIPs("subnet-empty")
+	ips, err := ipam.AllocatedIPs(t.Context(), "subnet-empty")
 	require.NoError(t, err)
 	assert.Nil(t, ips)
 
 	// Allocate some IPs
-	_, _ = ipam.AllocateIP("subnet-list", "10.0.7.0/24", PurposeENIPrimary, "eni-list")
-	_, _ = ipam.AllocateIP("subnet-list", "10.0.7.0/24", PurposeENIPrimary, "eni-list")
+	_, _ = ipam.AllocateIP(t.Context(), "subnet-list", "10.0.7.0/24", PurposeENIPrimary, "eni-list")
+	_, _ = ipam.AllocateIP(t.Context(), "subnet-list", "10.0.7.0/24", PurposeENIPrimary, "eni-list")
 
-	ips, err = ipam.AllocatedIPs("subnet-list")
+	ips, err = ipam.AllocatedIPs(t.Context(), "subnet-list")
 	require.NoError(t, err)
 	assert.Len(t, ips, 2)
 	assert.Equal(t, "10.0.7.4", ips[0].IP)
@@ -145,16 +146,16 @@ func TestIPAM_AllocatedIPs(t *testing.T) {
 func TestIPAM_MultipleSubnets(t *testing.T) {
 	ipam := setupTestIPAM(t)
 
-	ip1, err := ipam.AllocateIP("subnet-a", "10.0.10.0/24", PurposeENIPrimary, "eni-a")
+	ip1, err := ipam.AllocateIP(t.Context(), "subnet-a", "10.0.10.0/24", PurposeENIPrimary, "eni-a")
 	require.NoError(t, err)
 	assert.Equal(t, "10.0.10.4", ip1)
 
-	ip2, err := ipam.AllocateIP("subnet-b", "10.0.20.0/24", PurposeENIPrimary, "eni-b")
+	ip2, err := ipam.AllocateIP(t.Context(), "subnet-b", "10.0.20.0/24", PurposeENIPrimary, "eni-b")
 	require.NoError(t, err)
 	assert.Equal(t, "10.0.20.4", ip2)
 
 	// Each subnet tracks independently
-	ip3, err := ipam.AllocateIP("subnet-a", "10.0.10.0/24", PurposeENIPrimary, "eni-a")
+	ip3, err := ipam.AllocateIP(t.Context(), "subnet-a", "10.0.10.0/24", PurposeENIPrimary, "eni-a")
 	require.NoError(t, err)
 	assert.Equal(t, "10.0.10.5", ip3)
 }
@@ -163,15 +164,16 @@ func TestIPAM_LargerSubnet(t *testing.T) {
 	ipam := setupTestIPAM(t)
 
 	// /20 subnet — first allocable is still .4
-	ip, err := ipam.AllocateIP("subnet-big", "172.16.0.0/20", PurposeENIPrimary, "eni-big")
+	ip, err := ipam.AllocateIP(t.Context(), "subnet-big", "172.16.0.0/20", PurposeENIPrimary, "eni-big")
 	require.NoError(t, err)
 	assert.Equal(t, "172.16.0.4", ip)
 }
 
 func TestIPAM_NewIPAMWithKV(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 
-	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{
+	kv, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{
 		Bucket: "test-ipam-kv",
 	})
 	require.NoError(t, err)
@@ -180,14 +182,14 @@ func TestIPAM_NewIPAMWithKV(t *testing.T) {
 	require.NotNil(t, ipam)
 
 	// Should work the same as regular IPAM
-	ip, err := ipam.AllocateIP("subnet-kv", "10.0.1.0/24", PurposeENIPrimary, "eni-kv")
+	ip, err := ipam.AllocateIP(t.Context(), "subnet-kv", "10.0.1.0/24", PurposeENIPrimary, "eni-kv")
 	require.NoError(t, err)
 	assert.Equal(t, "10.0.1.4", ip)
 }
 
 func TestIPAM_InvalidCIDR(t *testing.T) {
 	ipam := setupTestIPAM(t)
-	_, err := ipam.AllocateIP("subnet-bad", "not-a-cidr", PurposeENIPrimary, "eni-bad")
+	_, err := ipam.AllocateIP(t.Context(), "subnet-bad", "not-a-cidr", PurposeENIPrimary, "eni-bad")
 	assert.Error(t, err)
 	assert.Contains(t, err.Error(), "parse CIDR")
 }

spinifex/handlers/ec2/vpc/security_group.go

@@ -17,7 +17,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
 	"github.com/mulgadc/spinifex/spinifex/utils"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // sgIDRegex must stay in lockstep with utils.GenerateResourceID("sg").
@@ -123,15 +123,15 @@ func (s *VPCServiceImpl) CreateSecurityGroup(ctx context.Context, input *ec2.Cre
 		return nil, errors.New(awserrors.ErrorInvalidGroupReserved)
 	}
 
-	if err := s.requireVPCExists(accountID, vpcId); err != nil {
+	if err := s.requireVPCExists(ctx, accountID, vpcId); err != nil {
 		return nil, err
 	}
 
 	// Check for duplicate group name in the same VPC and enforce the per-VPC
 	// SG quota in the same bucket walk.
 	prefix := accountID + "."
-	sgKeys, err := s.sgKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	sgKeys, err := s.sgKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 	sgsInVPC := 0
@@ -142,11 +142,11 @@ func (s *VPCServiceImpl) CreateSecurityGroup(ctx context.Context, input *ec2.Cre
 		if !strings.HasPrefix(k, prefix) {
 			continue
 		}
-		entry, err := s.sgKV.Get(k)
+		entry, err := s.sgKV.Get(ctx, k)
 		if err != nil {
 			// Fail closed — a transient read error must not let a duplicate
 			// SG name slip past, nor undercount the per-VPC quota.
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
 			}
 			slog.WarnContext(ctx, "CreateSecurityGroup: SG read failed", "key", k, "err", err)
@@ -196,7 +196,7 @@ func (s *VPCServiceImpl) CreateSecurityGroup(ctx context.Context, input *ec2.Cre
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal security group record: %w", err)
 	}
-	if _, err := s.sgKV.Put(utils.AccountKey(accountID, groupId), data); err != nil {
+	if _, err := s.sgKV.Put(ctx, utils.AccountKey(accountID, groupId), data); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -226,12 +226,12 @@ func (s *VPCServiceImpl) DeleteSecurityGroup(ctx context.Context, input *ec2.Del
 	groupId := *input.GroupId
 	key := utils.AccountKey(accountID, groupId)
 
-	entry, err := s.sgKV.Get(key)
+	entry, err := s.sgKV.Get(ctx, key)
 	if err != nil {
 		// AWS-faithful: an absent security group is InvalidGroup.NotFound, not
 		// success. Destroy orchestration tolerates it via awserrors.IsNotFound;
 		// a transient read error stays a server error.
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return nil, errors.New(awserrors.ErrorInvalidGroupNotFound)
 		}
 		return nil, errors.New(awserrors.ErrorServerInternal)
@@ -250,7 +250,7 @@ func (s *VPCServiceImpl) DeleteSecurityGroup(ctx context.Context, input *ec2.Del
 		return nil, err
 	}
 
-	if err := s.sgKV.Delete(key); err != nil {
+	if err := s.sgKV.Delete(ctx, key); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -270,12 +270,12 @@ func (s *VPCServiceImpl) DeleteSecurityGroup(ctx context.Context, input *ec2.Del
 // validateSGRuleReferences returns InvalidGroup.NotFound if any SourceSG in
 // rules is missing or belongs to a different VPC (cross-VPC refs not
 // supported; AWS uses the same error for both cases).
-func (s *VPCServiceImpl) validateSGRuleReferences(accountID, ownerVpcId string, rules []SGRule) error {
+func (s *VPCServiceImpl) validateSGRuleReferences(ctx context.Context, accountID, ownerVpcId string, rules []SGRule) error {
 	for _, r := range rules {
 		if r.SourceSG == "" {
 			continue
 		}
-		entry, err := s.sgKV.Get(utils.AccountKey(accountID, r.SourceSG))
+		entry, err := s.sgKV.Get(ctx, utils.AccountKey(accountID, r.SourceSG))
 		if err != nil {
 			return errors.New(awserrors.ErrorInvalidGroupNotFound)
 		}
@@ -296,15 +296,15 @@ func (s *VPCServiceImpl) validateSGRuleReferences(accountID, ownerVpcId string,
 func (s *VPCServiceImpl) checkSGDependencies(ctx context.Context, accountID, groupId string) error {
 	prefix := accountID + "."
 
-	eniKeys, err := s.eniKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	eniKeys, err := s.eniKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return errors.New(awserrors.ErrorServerInternal)
 	}
 	for _, k := range eniKeys {
 		if k == utils.VersionKey || !strings.HasPrefix(k, prefix) {
 			continue
 		}
-		entry, err := s.eniKV.Get(k)
+		entry, err := s.eniKV.Get(ctx, k)
 		if err != nil {
 			// Fail closed — a transient read error must not let us delete an
 			// SG that's actually still attached.
@@ -321,15 +321,15 @@ func (s *VPCServiceImpl) checkSGDependencies(ctx context.Context, accountID, gro
 		}
 	}
 
-	sgKeys, err := s.sgKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	sgKeys, err := s.sgKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return errors.New(awserrors.ErrorServerInternal)
 	}
 	for _, k := range sgKeys {
 		if k == utils.VersionKey || !strings.HasPrefix(k, prefix) {
 			continue
 		}
-		entry, err := s.sgKV.Get(k)
+		entry, err := s.sgKV.Get(ctx, k)
 		if err != nil {
 			slog.WarnContext(ctx, "checkSGDependencies: SG read failed", "key", k, "err", err)
 			return errors.New(awserrors.ErrorServerInternal)
@@ -383,8 +383,8 @@ func (s *VPCServiceImpl) DescribeSecurityGroups(ctx context.Context, input *ec2.
 	}
 
 	prefix := accountID + "."
-	keys, err := s.sgKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	keys, err := s.sgKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -396,7 +396,7 @@ func (s *VPCServiceImpl) DescribeSecurityGroups(ctx context.Context, input *ec2.
 			continue
 		}
 
-		entry, err := s.sgKV.Get(key)
+		entry, err := s.sgKV.Get(ctx, key)
 		if err != nil {
 			slog.WarnContext(ctx, "Failed to get security group record", "key", key, "error", err)
 			continue
@@ -519,8 +519,8 @@ func (s *VPCServiceImpl) DescribeSecurityGroupRules(ctx context.Context, input *
 	}
 
 	prefix := accountID + "."
-	keys, err := s.sgKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	keys, err := s.sgKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -531,7 +531,7 @@ func (s *VPCServiceImpl) DescribeSecurityGroupRules(ctx context.Context, input *
 		if key == utils.VersionKey || !strings.HasPrefix(key, prefix) {
 			continue
 		}
-		entry, err := s.sgKV.Get(key)
+		entry, err := s.sgKV.Get(ctx, key)
 		if err != nil {
 			slog.ErrorContext(ctx, "DescribeSecurityGroupRules: SG read failed", "key", key, "err", err)
 			return nil, errors.New(awserrors.ErrorServerInternal)
@@ -645,7 +645,7 @@ func (s *VPCServiceImpl) AuthorizeSecurityGroupIngress(ctx context.Context, inpu
 	groupId := *input.GroupId
 	key := utils.AccountKey(accountID, groupId)
 
-	entry, err := s.sgKV.Get(key)
+	entry, err := s.sgKV.Get(ctx, key)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidGroupNotFound)
 	}
@@ -655,7 +655,7 @@ func (s *VPCServiceImpl) AuthorizeSecurityGroupIngress(ctx context.Context, inpu
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
-	if err := s.requireVPCExists(accountID, record.VpcId); err != nil {
+	if err := s.requireVPCExists(ctx, accountID, record.VpcId); err != nil {
 		return nil, err
 	}
 
@@ -664,7 +664,7 @@ func (s *VPCServiceImpl) AuthorizeSecurityGroupIngress(ctx context.Context, inpu
 		slog.WarnContext(ctx, "AuthorizeSecurityGroupIngress: invalid rule", "groupId", groupId, "err", err)
 		return nil, errors.New(awserrors.ErrorInvalidParameterValue)
 	}
-	if err := s.validateSGRuleReferences(accountID, record.VpcId, newRules); err != nil {
+	if err := s.validateSGRuleReferences(ctx, accountID, record.VpcId, newRules); err != nil {
 		return nil, err
 	}
 	existing := make(map[string]struct{}, len(record.IngressRules))
@@ -685,7 +685,7 @@ func (s *VPCServiceImpl) AuthorizeSecurityGroupIngress(ctx context.Context, inpu
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal security group record: %w", err)
 	}
-	if _, err := s.sgKV.Update(key, data, entry.Revision()); err != nil {
+	if _, err := s.sgKV.Update(ctx, key, data, entry.Revision()); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -720,7 +720,7 @@ func (s *VPCServiceImpl) AuthorizeSecurityGroupEgress(ctx context.Context, input
 	groupId := *input.GroupId
 	key := utils.AccountKey(accountID, groupId)
 
-	entry, err := s.sgKV.Get(key)
+	entry, err := s.sgKV.Get(ctx, key)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidGroupNotFound)
 	}
@@ -730,7 +730,7 @@ func (s *VPCServiceImpl) AuthorizeSecurityGroupEgress(ctx context.Context, input
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
-	if err := s.requireVPCExists(accountID, record.VpcId); err != nil {
+	if err := s.requireVPCExists(ctx, accountID, record.VpcId); err != nil {
 		return nil, err
 	}
 
@@ -739,7 +739,7 @@ func (s *VPCServiceImpl) AuthorizeSecurityGroupEgress(ctx context.Context, input
 		slog.WarnContext(ctx, "AuthorizeSecurityGroupEgress: invalid rule", "groupId", groupId, "err", err)
 		return nil, errors.New(awserrors.ErrorInvalidParameterValue)
 	}
-	if err := s.validateSGRuleReferences(accountID, record.VpcId, newRules); err != nil {
+	if err := s.validateSGRuleReferences(ctx, accountID, record.VpcId, newRules); err != nil {
 		return nil, err
 	}
 	existing := make(map[string]struct{}, len(record.EgressRules))
@@ -760,7 +760,7 @@ func (s *VPCServiceImpl) AuthorizeSecurityGroupEgress(ctx context.Context, input
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal security group record: %w", err)
 	}
-	if _, err := s.sgKV.Update(key, data, entry.Revision()); err != nil {
+	if _, err := s.sgKV.Update(ctx, key, data, entry.Revision()); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -795,7 +795,7 @@ func (s *VPCServiceImpl) RevokeSecurityGroupIngress(ctx context.Context, input *
 	groupId := *input.GroupId
 	key := utils.AccountKey(accountID, groupId)
 
-	entry, err := s.sgKV.Get(key)
+	entry, err := s.sgKV.Get(ctx, key)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidGroupNotFound)
 	}
@@ -827,7 +827,7 @@ func (s *VPCServiceImpl) RevokeSecurityGroupIngress(ctx context.Context, input *
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal security group record: %w", err)
 	}
-	if _, err := s.sgKV.Update(key, data, entry.Revision()); err != nil {
+	if _, err := s.sgKV.Update(ctx, key, data, entry.Revision()); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -857,7 +857,7 @@ func (s *VPCServiceImpl) RevokeSecurityGroupEgress(ctx context.Context, input *e
 	groupId := *input.GroupId
 	key := utils.AccountKey(accountID, groupId)
 
-	entry, err := s.sgKV.Get(key)
+	entry, err := s.sgKV.Get(ctx, key)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidGroupNotFound)
 	}
@@ -889,7 +889,7 @@ func (s *VPCServiceImpl) RevokeSecurityGroupEgress(ctx context.Context, input *e
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal security group record: %w", err)
 	}
-	if _, err := s.sgKV.Update(key, data, entry.Revision()); err != nil {
+	if _, err := s.sgKV.Update(ctx, key, data, entry.Revision()); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -1171,7 +1171,7 @@ func (s *VPCServiceImpl) createDefaultSecurityGroupInternal(ctx context.Context,
 	if err != nil {
 		return "", fmt.Errorf("marshal default security group: %w", err)
 	}
-	if _, err := s.sgKV.Put(utils.AccountKey(accountID, groupId), data); err != nil {
+	if _, err := s.sgKV.Put(ctx, utils.AccountKey(accountID, groupId), data); err != nil {
 		return "", fmt.Errorf("store default security group: %w", err)
 	}
 
@@ -1193,7 +1193,7 @@ func (s *VPCServiceImpl) createDefaultSecurityGroupInternal(ctx context.Context,
 // per-VPC default SG. Surfaces vpcd tear-down failures to the caller.
 func (s *VPCServiceImpl) deleteSecurityGroupInternal(ctx context.Context, accountID, groupId string) error {
 	key := utils.AccountKey(accountID, groupId)
-	entry, err := s.sgKV.Get(key)
+	entry, err := s.sgKV.Get(ctx, key)
 	if err != nil {
 		return fmt.Errorf("read default security group: %w", err)
 	}
@@ -1201,7 +1201,7 @@ func (s *VPCServiceImpl) deleteSecurityGroupInternal(ctx context.Context, accoun
 	if err := json.Unmarshal(entry.Value(), &record); err != nil {
 		return fmt.Errorf("unmarshal default security group: %w", err)
 	}
-	if err := s.sgKV.Delete(key); err != nil {
+	if err := s.sgKV.Delete(ctx, key); err != nil {
 		return fmt.Errorf("delete default security group: %w", err)
 	}
 	return s.requestSGEvent("vpc.delete-sg", SGEvent{
@@ -1214,10 +1214,14 @@ func (s *VPCServiceImpl) deleteSecurityGroupInternal(ctx context.Context, accoun
 // IsDefault=true and the given VPC. Returns "" if none found (only happens
 // when CreateVpc failed mid-flow and left the VPC without a default SG).
 func (s *VPCServiceImpl) FindDefaultSGForVPC(accountID, vpcId string) (string, error) {
+	return s.findDefaultSGForVPC(context.Background(), accountID, vpcId)
+}
+
+func (s *VPCServiceImpl) findDefaultSGForVPC(ctx context.Context, accountID, vpcId string) (string, error) {
 	prefix := accountID + "."
-	keys, err := s.sgKV.Keys()
+	keys, err := s.sgKV.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return "", nil
 		}
 		return "", err
@@ -1226,7 +1230,7 @@ func (s *VPCServiceImpl) FindDefaultSGForVPC(accountID, vpcId string) (string, e
 		if k == utils.VersionKey || !strings.HasPrefix(k, prefix) {
 			continue
 		}
-		entry, err := s.sgKV.Get(k)
+		entry, err := s.sgKV.Get(ctx, k)
 		if err != nil {
 			continue
 		}

spinifex/handlers/ec2/vpc/security_group_test.go

@@ -1118,14 +1118,14 @@ func TestDescribeSecurityGroups_FilterNoResults(t *testing.T) {
 // IsDefault=true (production code keys off IsDefault, not GroupName).
 func findDefaultSGInVPC(t *testing.T, svc *VPCServiceImpl, vpcID string) string {
 	t.Helper()
-	keys, err := svc.sgKV.Keys()
+	keys, err := svc.sgKV.Keys(t.Context())
 	require.NoError(t, err)
 	var matches []string
 	for _, k := range keys {
 		if k == utils.VersionKey {
 			continue
 		}
-		entry, err := svc.sgKV.Get(k)
+		entry, err := svc.sgKV.Get(t.Context(), k)
 		require.NoError(t, err)
 		var rec SecurityGroupRecord
 		require.NoError(t, json.Unmarshal(entry.Value(), &rec))
@@ -1611,7 +1611,7 @@ func TestFindDefaultSGForVPC_SkipsMalformedRecord(t *testing.T) {
 	svc := setupTestVPCService(t)
 	vpcID := createTestVPC(t, svc, "10.0.0.0/16")
 
-	_, err := svc.sgKV.Put(testAccountID+".sg-corrupt0000000", []byte("{not json"))
+	_, err := svc.sgKV.Put(t.Context(), testAccountID+".sg-corrupt0000000", []byte("{not json"))
 	require.NoError(t, err)
 
 	// CreateVpc above already provisioned the default SG; FindDefaultSGForVPC
@@ -1651,7 +1651,7 @@ func TestDeleteSecurityGroup_FailsClosedOnCorruptENI(t *testing.T) {
 
 	// Inject a malformed ENI record into the bucket; checkSGDependencies must
 	// reject the delete rather than silently treating it as "no dependent ENI".
-	_, err := svc.eniKV.Put(testAccountID+".eni-corrupt", []byte("{not json"))
+	_, err := svc.eniKV.Put(t.Context(), testAccountID+".eni-corrupt", []byte("{not json"))
 	require.NoError(t, err)
 
 	_, err = svc.DeleteSecurityGroup(context.Background(), &ec2.DeleteSecurityGroupInput{
@@ -1669,7 +1669,7 @@ func seedAvailableVPC(t *testing.T, svc *VPCServiceImpl, vpcID string) {
 	rec := VPCRecord{VpcId: vpcID, CidrBlock: "10.0.0.0/16", State: "available"}
 	data, err := json.Marshal(rec)
 	require.NoError(t, err)
-	_, err = svc.vpcKV.Put(utils.AccountKey(testAccountID, vpcID), data)
+	_, err = svc.vpcKV.Put(t.Context(), utils.AccountKey(testAccountID, vpcID), data)
 	require.NoError(t, err)
 }
 
@@ -1894,7 +1894,7 @@ func TestAuthorizeSecurityGroupIngress_AssignsRuleIDs(t *testing.T) {
 	}, testAccountID)
 	require.NoError(t, err)
 
-	entry, err := svc.sgKV.Get(utils.AccountKey(testAccountID, sgID))
+	entry, err := svc.sgKV.Get(t.Context(), utils.AccountKey(testAccountID, sgID))
 	require.NoError(t, err)
 	var record SecurityGroupRecord
 	require.NoError(t, json.Unmarshal(entry.Value(), &record))

spinifex/handlers/ec2/vpc/service_impl.go

@@ -16,9 +16,11 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Ensure VPCServiceImpl implements VPCService.
@@ -74,17 +76,17 @@ type SubnetRecord struct {
 type VPCServiceImpl struct {
 	config   *config.Config
 	natsConn *nats.Conn
-	vpcKV    nats.KeyValue
-	subnetKV nats.KeyValue
-	vniKV    nats.KeyValue
-	eniKV    nats.KeyValue
-	sgKV     nats.KeyValue
-	rtbKV    nats.KeyValue // route table bucket for auto-creating main route table
+	vpcKV    jetstream.KeyValue
+	subnetKV jetstream.KeyValue
+	vniKV    jetstream.KeyValue
+	eniKV    jetstream.KeyValue
+	sgKV     jetstream.KeyValue
+	rtbKV    jetstream.KeyValue // route table bucket for auto-creating main route table
 	ipam     *IPAM
 
 	// Optional: injected after construction for public IP cleanup in DeleteNetworkInterface.
 	externalIPAM *ExternalIPAM
-	eipKV        nats.KeyValue
+	eipKV        jetstream.KeyValue
 
 	// disableDefaultPublicIP seeds default subnets with MapPublicIpOnLaunch=false
 	// (non-pool external modes have no public IPs to assign). Zero value keeps
@@ -100,7 +102,7 @@ func (s *VPCServiceImpl) SetDefaultPublicIPMapping(enabled bool) {
 
 // SetExternalIPAM injects external IPAM and EIP KV store references so that
 // DeleteNetworkInterface can release auto-assigned public IPs and NAT rules.
-func (s *VPCServiceImpl) SetExternalIPAM(ipam *ExternalIPAM, eipKV nats.KeyValue) {
+func (s *VPCServiceImpl) SetExternalIPAM(ipam *ExternalIPAM, eipKV jetstream.KeyValue) {
 	s.externalIPAM = ipam
 	s.eipKV = eipKV
 }
@@ -116,58 +118,58 @@ func (s *VPCServiceImpl) localAZ() string {
 }
 
 // NewVPCServiceImplWithNATS creates a VPC service with NATS JetStream for persistence.
-func NewVPCServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*VPCServiceImpl, error) {
-	js, err := natsConn.JetStream()
+func NewVPCServiceImplWithNATS(ctx context.Context, cfg *config.Config, natsConn *nats.Conn) (*VPCServiceImpl, error) {
+	js, err := jetstream.New(natsConn)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get JetStream context: %w", err)
+		return nil, fmt.Errorf("create JetStream client: %w", err)
 	}
 
-	vpcKV, err := utils.GetOrCreateKVBucket(js, KVBucketVPCs, 10)
+	vpcKV, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketVPCs, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketVPCs, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketVPCs, vpcKV, KVBucketVPCsVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketVPCs, vpcKV, KVBucketVPCsVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketVPCs, err)
 	}
 
-	subnetKV, err := utils.GetOrCreateKVBucket(js, KVBucketSubnets, 10)
+	subnetKV, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketSubnets, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketSubnets, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketSubnets, subnetKV, KVBucketSubnetsVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketSubnets, subnetKV, KVBucketSubnetsVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketSubnets, err)
 	}
 
-	vniKV, err := utils.GetOrCreateKVBucket(js, KVBucketVNICounter, 10)
+	vniKV, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketVNICounter, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketVNICounter, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketVNICounter, vniKV, KVBucketVNICounterVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketVNICounter, vniKV, KVBucketVNICounterVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketVNICounter, err)
 	}
 
-	eniKV, err := utils.GetOrCreateKVBucket(js, KVBucketENIs, 10)
+	eniKV, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketENIs, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketENIs, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketENIs, eniKV, KVBucketENIsVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketENIs, eniKV, KVBucketENIsVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketENIs, err)
 	}
 
-	sgKV, err := utils.GetOrCreateKVBucket(js, KVBucketSecurityGroups, 10)
+	sgKV, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketSecurityGroups, 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketSecurityGroups, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketSecurityGroups, sgKV, KVBucketSecurityGroupsVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketSecurityGroups, sgKV, KVBucketSecurityGroupsVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketSecurityGroups, err)
 	}
 
-	rtbKV, err := utils.GetOrCreateKVBucket(js, "spinifex-vpc-route-tables", 10)
+	rtbKV, err := kvutil.GetOrCreateBucket(ctx, js, "spinifex-vpc-route-tables", 10)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create KV bucket spinifex-vpc-route-tables: %w", err)
 	}
 
-	ipam, err := NewIPAM(js)
+	ipam, err := NewIPAM(ctx, js)
 	if err != nil {
 		return nil, fmt.Errorf("failed to initialize IPAM: %w", err)
 	}
@@ -193,17 +195,17 @@ func NewVPCServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*VPCSer
 }
 
 // nextVNI allocates the next VNI using atomic increment on the NATS KV counter.
-func (s *VPCServiceImpl) nextVNI() (int64, error) {
-	entry, err := s.vniKV.Get(vniCounterKey)
+func (s *VPCServiceImpl) nextVNI(ctx context.Context) (int64, error) {
+	entry, err := s.vniKV.Get(ctx, vniCounterKey)
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			// First VNI allocation — initialize counter
 			vni := int64(vniStart)
 			data, marshalErr := json.Marshal(vni + 1)
 			if marshalErr != nil {
 				return 0, fmt.Errorf("failed to marshal VNI counter: %w", marshalErr)
 			}
-			if _, err := s.vniKV.Create(vniCounterKey, data); err != nil {
+			if _, err := s.vniKV.Create(ctx, vniCounterKey, data); err != nil {
 				return 0, fmt.Errorf("failed to initialize VNI counter: %w", err)
 			}
 			return vni, nil
@@ -221,7 +223,7 @@ func (s *VPCServiceImpl) nextVNI() (int64, error) {
 	if marshalErr != nil {
 		return 0, fmt.Errorf("failed to marshal VNI counter: %w", marshalErr)
 	}
-	if _, err := s.vniKV.Update(vniCounterKey, data, entry.Revision()); err != nil {
+	if _, err := s.vniKV.Update(ctx, vniCounterKey, data, entry.Revision()); err != nil {
 		return 0, fmt.Errorf("failed to update VNI counter (CAS conflict): %w", err)
 	}
 
@@ -247,7 +249,7 @@ func (s *VPCServiceImpl) CreateVpc(ctx context.Context, input *ec2.CreateVpcInpu
 	}
 
 	// Allocate VNI for overlay network
-	vni, err := s.nextVNI()
+	vni, err := s.nextVNI(ctx)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
@@ -271,7 +273,7 @@ func (s *VPCServiceImpl) CreateVpc(ctx context.Context, input *ec2.CreateVpcInpu
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal VPC record: %w", err)
 	}
-	if _, err := s.vpcKV.Put(utils.AccountKey(accountID, vpcID), data); err != nil {
+	if _, err := s.vpcKV.Put(ctx, utils.AccountKey(accountID, vpcID), data); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -301,8 +303,8 @@ func (s *VPCServiceImpl) CreateVpc(ctx context.Context, input *ec2.CreateVpcInpu
 
 // requireVPCExists returns InvalidVpcID.NotFound if the VPC doesn't exist for
 // this account.
-func (s *VPCServiceImpl) requireVPCExists(accountID, vpcId string) error {
-	if _, err := s.vpcKV.Get(utils.AccountKey(accountID, vpcId)); err != nil {
+func (s *VPCServiceImpl) requireVPCExists(ctx context.Context, accountID, vpcId string) error {
+	if _, err := s.vpcKV.Get(ctx, utils.AccountKey(accountID, vpcId)); err != nil {
 		return errors.New(awserrors.ErrorInvalidVpcIDNotFound)
 	}
 	return nil
@@ -317,11 +319,11 @@ func (s *VPCServiceImpl) DeleteVpc(ctx context.Context, input *ec2.DeleteVpcInpu
 	vpcID := *input.VpcId
 	key := utils.AccountKey(accountID, vpcID)
 
-	if _, err := s.vpcKV.Get(key); err != nil {
+	if _, err := s.vpcKV.Get(ctx, key); err != nil {
 		// AWS-faithful: an absent VPC is NotFound (the tofu/SDK provider
 		// tolerates it on destroy); destroy orchestration tolerates it too.
 		// A transient read error stays a server error.
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return nil, errors.New(awserrors.ErrorInvalidVpcIDNotFound)
 		}
 		return nil, errors.New(awserrors.ErrorServerInternal)
@@ -329,8 +331,8 @@ func (s *VPCServiceImpl) DeleteVpc(ctx context.Context, input *ec2.DeleteVpcInpu
 
 	// Check for dependent subnets owned by this account
 	prefix := accountID + "."
-	subnetKeys, err := s.subnetKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	subnetKeys, err := s.subnetKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 	for _, k := range subnetKeys {
@@ -340,13 +342,13 @@ func (s *VPCServiceImpl) DeleteVpc(ctx context.Context, input *ec2.DeleteVpcInpu
 		if !strings.HasPrefix(k, prefix) {
 			continue
 		}
-		entry, err := s.subnetKV.Get(k)
+		entry, err := s.subnetKV.Get(ctx, k)
 		if err != nil {
 			// ErrKeyNotFound means the subnet was deleted between Keys() and
 			// Get() — fine to skip. Any other error is fail-closed: a
 			// transient read error must not let DeleteVpc bypass a subnet
 			// dependency it can't see.
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
 			}
 			slog.WarnContext(ctx, "DeleteVpc: subnet read failed", "key", k, "err", err)
@@ -365,21 +367,21 @@ func (s *VPCServiceImpl) DeleteVpc(ctx context.Context, input *ec2.DeleteVpcInpu
 	// Reject if any non-default SG remains in this VPC; the cascade only
 	// auto-deletes the default SG (matches AWS DeleteVpc semantics).
 	defaultSGId := ""
-	sgKeys, err := s.sgKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	sgKeys, err := s.sgKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 	for _, k := range sgKeys {
 		if k == utils.VersionKey || !strings.HasPrefix(k, prefix) {
 			continue
 		}
-		entry, err := s.sgKV.Get(k)
+		entry, err := s.sgKV.Get(ctx, k)
 		if err != nil {
 			// ErrKeyNotFound means the SG was deleted between Keys() and
 			// Get() — fine to skip. Any other error is fail-closed: a
 			// transient read error must not let DeleteVpc orphan a
 			// non-default SG it can't see.
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
 			}
 			slog.WarnContext(ctx, "DeleteVpc: SG read failed", "key", k, "err", err)
@@ -411,20 +413,20 @@ func (s *VPCServiceImpl) DeleteVpc(ctx context.Context, input *ec2.DeleteVpcInpu
 	// Reap the VPC's main route table: CreateVpc auto-creates it and
 	// DeleteRouteTable refuses to delete a main RT, so nothing else reclaims it
 	// — without this rtbKV leaks one orphaned main RT per deleted VPC.
-	rtbID, err := s.findMainRouteTableID(accountID, vpcID)
+	rtbID, err := s.findMainRouteTableID(ctx, accountID, vpcID)
 	if err != nil {
 		slog.ErrorContext(ctx, "DeleteVpc: main route table lookup failed", "vpcId", vpcID, "err", err)
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 	if rtbID != "" {
-		if err := s.rtbKV.Delete(utils.AccountKey(accountID, rtbID)); err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
+		if err := s.rtbKV.Delete(ctx, utils.AccountKey(accountID, rtbID)); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
 			slog.ErrorContext(ctx, "DeleteVpc: main route table reap failed", "vpcId", vpcID, "routeTableId", rtbID, "err", err)
 			return nil, errors.New(awserrors.ErrorServerInternal)
 		}
 		slog.InfoContext(ctx, "DeleteVpc: reaped main route table", "vpcId", vpcID, "routeTableId", rtbID)
 	}
 
-	if err := s.vpcKV.Delete(key); err != nil {
+	if err := s.vpcKV.Delete(ctx, key); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -463,8 +465,8 @@ func (s *VPCServiceImpl) DescribeVpcs(ctx context.Context, input *ec2.DescribeVp
 	}
 
 	prefix := accountID + "."
-	keys, err := s.vpcKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	keys, err := s.vpcKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -476,7 +478,7 @@ func (s *VPCServiceImpl) DescribeVpcs(ctx context.Context, input *ec2.DescribeVp
 			continue
 		}
 
-		entry, err := s.vpcKV.Get(key)
+		entry, err := s.vpcKV.Get(ctx, key)
 		if err != nil {
 			slog.WarnContext(ctx, "Failed to get VPC record", "key", key, "error", err)
 			continue
@@ -533,7 +535,7 @@ func (s *VPCServiceImpl) CreateSubnet(ctx context.Context, input *ec2.CreateSubn
 	vpcID := *input.VpcId
 
 	// Verify VPC exists and belongs to this account
-	vpcEntry, err := s.vpcKV.Get(utils.AccountKey(accountID, vpcID))
+	vpcEntry, err := s.vpcKV.Get(ctx, utils.AccountKey(accountID, vpcID))
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidVpcIDNotFound)
 	}
@@ -566,8 +568,8 @@ func (s *VPCServiceImpl) CreateSubnet(ctx context.Context, input *ec2.CreateSubn
 
 	// Check for CIDR conflicts with existing subnets in this VPC (same account)
 	prefix := accountID + "."
-	subnetKeys, err := s.subnetKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	subnetKeys, err := s.subnetKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -578,7 +580,7 @@ func (s *VPCServiceImpl) CreateSubnet(ctx context.Context, input *ec2.CreateSubn
 		if !strings.HasPrefix(k, prefix) {
 			continue
 		}
-		entry, err := s.subnetKV.Get(k)
+		entry, err := s.subnetKV.Get(ctx, k)
 		if err != nil {
 			continue
 		}
@@ -627,7 +629,7 @@ func (s *VPCServiceImpl) CreateSubnet(ctx context.Context, input *ec2.CreateSubn
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal subnet record: %w", err)
 	}
-	if _, err := s.subnetKV.Put(utils.AccountKey(accountID, subnetID), data); err != nil {
+	if _, err := s.subnetKV.Put(ctx, utils.AccountKey(accountID, subnetID), data); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -651,12 +653,12 @@ func (s *VPCServiceImpl) DeleteSubnet(ctx context.Context, input *ec2.DeleteSubn
 	key := utils.AccountKey(accountID, subnetID)
 
 	// Read subnet record before deletion (needed for vpcd event)
-	subnetEntry, err := s.subnetKV.Get(key)
+	subnetEntry, err := s.subnetKV.Get(ctx, key)
 	if err != nil {
 		// AWS-faithful: an absent subnet is NotFound (provider tolerates it on
 		// destroy); destroy orchestration tolerates it too. A transient read
 		// error stays a server error.
-		if errors.Is(err, nats.ErrKeyNotFound) {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return nil, errors.New(awserrors.ErrorInvalidSubnetIDNotFound)
 		}
 		return nil, errors.New(awserrors.ErrorServerInternal)
@@ -671,7 +673,7 @@ func (s *VPCServiceImpl) DeleteSubnet(ctx context.Context, input *ec2.DeleteSubn
 		return nil, err
 	}
 
-	if err := s.subnetKV.Delete(key); err != nil {
+	if err := s.subnetKV.Delete(ctx, key); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -687,9 +689,9 @@ func (s *VPCServiceImpl) DeleteSubnet(ctx context.Context, input *ec2.DeleteSubn
 // subnet is a live attachment (rule #3). Fail-closed on a KV read error so a
 // transient fault never lets a delete orphan a port.
 func (s *VPCServiceImpl) checkSubnetResidents(ctx context.Context, accountID, subnetID string) error {
-	keys, err := s.eniKV.Keys()
+	keys, err := s.eniKV.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil
 		}
 		slog.ErrorContext(ctx, "DeleteSubnet: ENI scan failed, blocking delete to avoid orphaning ports", "subnetId", subnetID, "err", err)
@@ -700,9 +702,9 @@ func (s *VPCServiceImpl) checkSubnetResidents(ctx context.Context, accountID, su
 		if key == utils.VersionKey || !strings.HasPrefix(key, prefix) {
 			continue
 		}
-		entry, err := s.eniKV.Get(key)
+		entry, err := s.eniKV.Get(ctx, key)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
 			}
 			return errors.New(awserrors.ErrorServerInternal)
@@ -736,8 +738,8 @@ func (s *VPCServiceImpl) DescribeSubnets(ctx context.Context, input *ec2.Describ
 	}
 
 	prefix := accountID + "."
-	keys, err := s.subnetKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	keys, err := s.subnetKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -749,7 +751,7 @@ func (s *VPCServiceImpl) DescribeSubnets(ctx context.Context, input *ec2.Describ
 			continue
 		}
 
-		entry, err := s.subnetKV.Get(key)
+		entry, err := s.subnetKV.Get(ctx, key)
 		if err != nil {
 			slog.WarnContext(ctx, "Failed to get subnet record", "key", key, "error", err)
 			continue
@@ -815,7 +817,7 @@ func (s *VPCServiceImpl) ApplyRecordTags(input *ec2.CreateTagsInput, accountID s
 		if res == nil {
 			continue
 		}
-		if err := s.updateRecordTags(accountID, *res, merge); err != nil {
+		if err := s.updateRecordTags(context.Background(), accountID, *res, merge); err != nil {
 			return err
 		}
 	}
@@ -834,7 +836,7 @@ func (s *VPCServiceImpl) RemoveRecordTags(input *ec2.DeleteTagsInput, accountID
 		if res == nil {
 			continue
 		}
-		if err := s.updateRecordTags(accountID, *res, remove); err != nil {
+		if err := s.updateRecordTags(context.Background(), accountID, *res, remove); err != nil {
 			return err
 		}
 	}
@@ -843,31 +845,31 @@ func (s *VPCServiceImpl) RemoveRecordTags(input *ec2.DeleteTagsInput, accountID
 
 // updateRecordTags applies mut to the tag map of the subnet-, vpc-, sg-, or
 // eni-scoped record identified by resourceID. Other resource ids are a no-op.
-func (s *VPCServiceImpl) updateRecordTags(accountID, resourceID string, mut func(map[string]string)) error {
+func (s *VPCServiceImpl) updateRecordTags(ctx context.Context, accountID, resourceID string, mut func(map[string]string)) error {
 	switch {
 	case strings.HasPrefix(resourceID, "subnet-"):
-		return utils.UpdateKVRecordTags(s.subnetKV, accountID, resourceID, func(r *SubnetRecord) {
+		return utils.UpdateKVRecordTags(ctx, s.subnetKV, accountID, resourceID, func(r *SubnetRecord) {
 			if r.Tags == nil {
 				r.Tags = map[string]string{}
 			}
 			mut(r.Tags)
 		})
 	case strings.HasPrefix(resourceID, "vpc-"):
-		return utils.UpdateKVRecordTags(s.vpcKV, accountID, resourceID, func(r *VPCRecord) {
+		return utils.UpdateKVRecordTags(ctx, s.vpcKV, accountID, resourceID, func(r *VPCRecord) {
 			if r.Tags == nil {
 				r.Tags = map[string]string{}
 			}
 			mut(r.Tags)
 		})
 	case strings.HasPrefix(resourceID, "sg-"):
-		return utils.UpdateKVRecordTags(s.sgKV, accountID, resourceID, func(r *SecurityGroupRecord) {
+		return utils.UpdateKVRecordTags(ctx, s.sgKV, accountID, resourceID, func(r *SecurityGroupRecord) {
 			if r.Tags == nil {
 				r.Tags = map[string]string{}
 			}
 			mut(r.Tags)
 		})
 	case strings.HasPrefix(resourceID, "eni-"):
-		return utils.UpdateKVRecordTags(s.eniKV, accountID, resourceID, func(r *ENIRecord) {
+		return utils.UpdateKVRecordTags(ctx, s.eniKV, accountID, resourceID, func(r *ENIRecord) {
 			if r.Tags == nil {
 				r.Tags = map[string]string{}
 			}
@@ -1004,7 +1006,7 @@ func (s *VPCServiceImpl) ModifySubnetAttribute(ctx context.Context, input *ec2.M
 	subnetID := *input.SubnetId
 	key := utils.AccountKey(accountID, subnetID)
 
-	entry, err := s.subnetKV.Get(key)
+	entry, err := s.subnetKV.Get(ctx, key)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidSubnetIDNotFound)
 	}
@@ -1022,7 +1024,7 @@ func (s *VPCServiceImpl) ModifySubnetAttribute(ctx context.Context, input *ec2.M
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal subnet record: %w", err)
 	}
-	if _, err := s.subnetKV.Update(key, data, entry.Revision()); err != nil {
+	if _, err := s.subnetKV.Update(ctx, key, data, entry.Revision()); err != nil {
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -1043,7 +1045,7 @@ func (s *VPCServiceImpl) ModifyVpcAttribute(ctx context.Context, input *ec2.Modi
 	vpcID := *input.VpcId
 	key := utils.AccountKey(accountID, vpcID)
 
-	entry, err := s.vpcKV.Get(key)
+	entry, err := s.vpcKV.Get(ctx, key)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidVpcIDNotFound)
 	}
@@ -1069,7 +1071,7 @@ func (s *VPCServiceImpl) ModifyVpcAttribute(ctx context.Context, input *ec2.Modi
 		slog.ErrorContext(ctx, "ModifyVpcAttribute: failed to marshal VPC record", "vpcId", vpcID, "accountID", accountID, "err", err)
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
-	if _, err := s.vpcKV.Update(key, data, entry.Revision()); err != nil {
+	if _, err := s.vpcKV.Update(ctx, key, data, entry.Revision()); err != nil {
 		slog.ErrorContext(ctx, "ModifyVpcAttribute: KV update failed", "vpcId", vpcID, "accountID", accountID, "err", err)
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
@@ -1091,7 +1093,7 @@ func (s *VPCServiceImpl) DescribeVpcAttribute(ctx context.Context, input *ec2.De
 	vpcID := *input.VpcId
 	key := utils.AccountKey(accountID, vpcID)
 
-	entry, err := s.vpcKV.Get(key)
+	entry, err := s.vpcKV.Get(ctx, key)
 	if err != nil {
 		return nil, errors.New(awserrors.ErrorInvalidVpcIDNotFound)
 	}
@@ -1145,14 +1147,18 @@ type BootstrapIDs struct {
 // EnsureDefaultVPC creates a default VPC and subnet if none exists for the
 // account. Safe to call multiple times — no-ops if already present.
 func (s *VPCServiceImpl) EnsureDefaultVPC(accountID string, bootstrap ...BootstrapIDs) (*DefaultVPCInfo, error) {
+	return s.ensureDefaultVPC(context.Background(), accountID, bootstrap...)
+}
+
+func (s *VPCServiceImpl) ensureDefaultVPC(ctx context.Context, accountID string, bootstrap ...BootstrapIDs) (*DefaultVPCInfo, error) {
 	if s.vpcKV == nil {
 		return nil, nil // No persistence, skip
 	}
 
 	// Check if a default VPC already exists for this account
 	prefix := accountID + "."
-	keys, err := s.vpcKV.Keys()
-	if err != nil && !errors.Is(err, nats.ErrNoKeysFound) {
+	keys, err := s.vpcKV.Keys(ctx)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, fmt.Errorf("list VPCs: %w", err)
 	}
 
@@ -1163,7 +1169,7 @@ func (s *VPCServiceImpl) EnsureDefaultVPC(accountID string, bootstrap ...Bootstr
 		if !strings.HasPrefix(key, prefix) {
 			continue
 		}
-		entry, err := s.vpcKV.Get(key)
+		entry, err := s.vpcKV.Get(ctx, key)
 		if err != nil {
 			continue
 		}
@@ -1174,7 +1180,7 @@ func (s *VPCServiceImpl) EnsureDefaultVPC(accountID string, bootstrap ...Bootstr
 		if record.IsDefault {
 			slog.Debug("Default VPC already exists", "vpcId", record.VpcId, "accountID", accountID)
 			// Look up the default subnet to return full info
-			defaultSubnet, _ := s.GetDefaultSubnet(accountID)
+			defaultSubnet, _ := s.getDefaultSubnet(ctx, accountID)
 			info := &DefaultVPCInfo{VpcId: record.VpcId, Cidr: record.CidrBlock}
 			if defaultSubnet != nil {
 				info.SubnetId = defaultSubnet.SubnetId
@@ -1185,7 +1191,7 @@ func (s *VPCServiceImpl) EnsureDefaultVPC(accountID string, bootstrap ...Bootstr
 	}
 
 	// Create default VPC — use bootstrap IDs if provided for consistency
-	vni, err := s.nextVNI()
+	vni, err := s.nextVNI(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("allocate VNI for default VPC: %w", err)
 	}
@@ -1211,7 +1217,7 @@ func (s *VPCServiceImpl) EnsureDefaultVPC(accountID string, bootstrap ...Bootstr
 	if err != nil {
 		return nil, fmt.Errorf("marshal default VPC: %w", err)
 	}
-	if _, err := s.vpcKV.Put(utils.AccountKey(accountID, vpcID), data); err != nil {
+	if _, err := s.vpcKV.Put(ctx, utils.AccountKey(accountID, vpcID), data); err != nil {
 		return nil, fmt.Errorf("store default VPC: %w", err)
 	}
 
@@ -1244,7 +1250,7 @@ func (s *VPCServiceImpl) EnsureDefaultVPC(accountID string, bootstrap ...Bootstr
 	if err != nil {
 		return nil, fmt.Errorf("marshal default subnet: %w", err)
 	}
-	if _, err := s.subnetKV.Put(utils.AccountKey(accountID, subnetID), data); err != nil {
+	if _, err := s.subnetKV.Put(ctx, utils.AccountKey(accountID, subnetID), data); err != nil {
 		return nil, fmt.Errorf("store default subnet: %w", err)
 	}
 
@@ -1252,7 +1258,7 @@ func (s *VPCServiceImpl) EnsureDefaultVPC(accountID string, bootstrap ...Bootstr
 
 	// Create main route table with local route (written directly to KV to avoid circular import)
 	if s.rtbKV != nil {
-		if err := s.createMainRouteTable(context.Background(), accountID, vpcID, DefaultVPCCidr); err != nil {
+		if err := s.createMainRouteTable(ctx, accountID, vpcID, DefaultVPCCidr); err != nil {
 			slog.Error("Failed to create main route table for VPC", "vpcId", vpcID, "err", err)
 		}
 	}
@@ -1261,7 +1267,7 @@ func (s *VPCServiceImpl) EnsureDefaultVPC(accountID string, bootstrap ...Bootstr
 	// before vpcd has subscribed to vpc.create-sg, so the synchronous round-trip
 	// will time out on first boot. The SG record is already in KV; vpcd's
 	// reconcile-sgs loop creates the OVN port group on its first scan.
-	if _, err := s.createDefaultSecurityGroupInternal(context.Background(), accountID, vpcID); err != nil {
+	if _, err := s.createDefaultSecurityGroupInternal(ctx, accountID, vpcID); err != nil {
 		slog.Warn("Default security group bootstrap deferred to vpcd reconciler",
 			"vpcId", vpcID, "accountID", accountID, "err", err)
 	}
@@ -1286,7 +1292,7 @@ func (s *VPCServiceImpl) EnsureDefaultVPC(accountID string, bootstrap ...Bootstr
 // table KV bucket. Avoids a circular import; idempotent via findMainRouteTableID
 // so concurrent EnsureDefaultVPC calls do not mint duplicate main RTs.
 func (s *VPCServiceImpl) createMainRouteTable(ctx context.Context, accountID, vpcID, vpcCidr string) error {
-	if existing, err := s.findMainRouteTableID(accountID, vpcID); err != nil {
+	if existing, err := s.findMainRouteTableID(ctx, accountID, vpcID); err != nil {
 		return fmt.Errorf("check existing main route table: %w", err)
 	} else if existing != "" {
 		slog.InfoContext(ctx, "Main route table already exists for VPC, skipping creation",
@@ -1340,7 +1346,7 @@ func (s *VPCServiceImpl) createMainRouteTable(ctx context.Context, accountID, vp
 	if err != nil {
 		return fmt.Errorf("marshal main route table: %w", err)
 	}
-	if _, err := s.rtbKV.Put(utils.AccountKey(accountID, rtbID), data); err != nil {
+	if _, err := s.rtbKV.Put(ctx, utils.AccountKey(accountID, rtbID), data); err != nil {
 		return fmt.Errorf("store main route table: %w", err)
 	}
 
@@ -1351,11 +1357,11 @@ func (s *VPCServiceImpl) createMainRouteTable(ctx context.Context, accountID, vp
 // findMainRouteTableID returns the rtb-ID of the main route table for vpcID
 // in accountID, or "" if none exists. Partial unmarshal — only the fields
 // needed to identify a main RT, so this stays cheap on hot KV scans.
-func (s *VPCServiceImpl) findMainRouteTableID(accountID, vpcID string) (string, error) {
+func (s *VPCServiceImpl) findMainRouteTableID(ctx context.Context, accountID, vpcID string) (string, error) {
 	prefix := accountID + "."
-	keys, err := s.rtbKV.Keys()
+	keys, err := s.rtbKV.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return "", nil
 		}
 		return "", fmt.Errorf("list rtb keys: %w", err)
@@ -1364,9 +1370,9 @@ func (s *VPCServiceImpl) findMainRouteTableID(accountID, vpcID string) (string,
 		if key == utils.VersionKey || !strings.HasPrefix(key, prefix) {
 			continue
 		}
-		entry, err := s.rtbKV.Get(key)
+		entry, err := s.rtbKV.Get(ctx, key)
 		if err != nil {
-			if errors.Is(err, nats.ErrKeyNotFound) {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
 			}
 			return "", fmt.Errorf("read rtb %s: %w", key, err)
@@ -1388,10 +1394,14 @@ func (s *VPCServiceImpl) findMainRouteTableID(accountID, vpcID string) (string,
 
 // GetDefaultSubnet returns the default subnet for RunInstances when no SubnetId is specified.
 func (s *VPCServiceImpl) GetDefaultSubnet(accountID string) (*SubnetRecord, error) {
+	return s.getDefaultSubnet(context.Background(), accountID)
+}
+
+func (s *VPCServiceImpl) getDefaultSubnet(ctx context.Context, accountID string) (*SubnetRecord, error) {
 	prefix := accountID + "."
-	keys, err := s.subnetKV.Keys()
+	keys, err := s.subnetKV.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, fmt.Errorf("no default subnet found")
 		}
 		return nil, fmt.Errorf("list subnets: %w", err)
@@ -1404,7 +1414,7 @@ func (s *VPCServiceImpl) GetDefaultSubnet(accountID string) (*SubnetRecord, erro
 		if !strings.HasPrefix(key, prefix) {
 			continue
 		}
-		entry, err := s.subnetKV.Get(key)
+		entry, err := s.subnetKV.Get(ctx, key)
 		if err != nil {
 			continue
 		}
@@ -1422,8 +1432,12 @@ func (s *VPCServiceImpl) GetDefaultSubnet(accountID string) (*SubnetRecord, erro
 
 // GetSubnet looks up a SubnetRecord by ID.
 func (s *VPCServiceImpl) GetSubnet(accountID, subnetId string) (*SubnetRecord, error) {
+	return s.getSubnet(context.Background(), accountID, subnetId)
+}
+
+func (s *VPCServiceImpl) getSubnet(ctx context.Context, accountID, subnetId string) (*SubnetRecord, error) {
 	key := utils.AccountKey(accountID, subnetId)
-	entry, err := s.subnetKV.Get(key)
+	entry, err := s.subnetKV.Get(ctx, key)
 	if err != nil {
 		return nil, fmt.Errorf("subnet %s not found: %w", subnetId, err)
 	}

spinifex/handlers/ec2/vpc/service_impl_test.go

@@ -21,7 +21,7 @@ func setupTestVPCServiceWithNC(t *testing.T) (*VPCServiceImpl, *nats.Conn) {
 	t.Helper()
 	_, nc, _ := testutil.StartTestJetStream(t)
 
-	svc, err := NewVPCServiceImplWithNATS(nil, nc)
+	svc, err := NewVPCServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 	testutil.StubVpcdSGResponder(t, nc)
 	return svc, nc
@@ -39,7 +39,7 @@ func setupTestVPCService(t *testing.T) *VPCServiceImpl {
 func setupTestVPCServiceWithFailingVpcd(t *testing.T, errMsg string) (*VPCServiceImpl, *nats.Conn) {
 	t.Helper()
 	_, nc, _ := testutil.StartTestJetStream(t)
-	svc, err := NewVPCServiceImplWithNATS(nil, nc)
+	svc, err := NewVPCServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 	testutil.StubVpcdSGFailingResponder(t, nc, errMsg)
 	return svc, nc
@@ -586,13 +586,13 @@ func TestCreateMainRouteTable_Idempotent(t *testing.T) {
 	svc := setupTestVPCService(t)
 	vpcID := createTestVPC(t, svc, "10.99.0.0/16") // CreateVpc auto-calls createMainRouteTable
 
-	firstID, err := svc.findMainRouteTableID(testAccountID, vpcID)
+	firstID, err := svc.findMainRouteTableID(t.Context(), testAccountID, vpcID)
 	require.NoError(t, err)
 	require.NotEmpty(t, firstID, "CreateVpc should have created a main route table")
 
 	require.NoError(t, svc.createMainRouteTable(context.Background(), testAccountID, vpcID, "10.99.0.0/16"))
 
-	secondID, err := svc.findMainRouteTableID(testAccountID, vpcID)
+	secondID, err := svc.findMainRouteTableID(t.Context(), testAccountID, vpcID)
 	require.NoError(t, err)
 	assert.Equal(t, firstID, secondID, "second call must be a no-op")
 
@@ -608,7 +608,7 @@ func TestDeleteVpc_ReapsMainRouteTable(t *testing.T) {
 	svc := setupTestVPCService(t)
 	vpcID := createTestVPC(t, svc, "10.77.0.0/16") // auto-creates the main RT
 
-	rtbID, err := svc.findMainRouteTableID(testAccountID, vpcID)
+	rtbID, err := svc.findMainRouteTableID(t.Context(), testAccountID, vpcID)
 	require.NoError(t, err)
 	require.NotEmpty(t, rtbID, "CreateVpc should have created a main route table")
 	require.Equal(t, 1, countMainRouteTablesForVPC(t, svc, vpcID))
@@ -616,7 +616,7 @@ func TestDeleteVpc_ReapsMainRouteTable(t *testing.T) {
 	_, err = svc.DeleteVpc(context.Background(), &ec2.DeleteVpcInput{VpcId: aws.String(vpcID)}, testAccountID)
 	require.NoError(t, err)
 
-	gone, err := svc.findMainRouteTableID(testAccountID, vpcID)
+	gone, err := svc.findMainRouteTableID(t.Context(), testAccountID, vpcID)
 	require.NoError(t, err)
 	assert.Empty(t, gone, "DeleteVpc must reap the main route table, not leak it")
 	assert.Equal(t, 0, countMainRouteTablesForVPC(t, svc, vpcID))
@@ -626,13 +626,13 @@ func TestDeleteVpc_ReapsMainRouteTable(t *testing.T) {
 // state in tests. Returns the number of IsMain=true records for vpcID.
 func countMainRouteTablesForVPC(t *testing.T, svc *VPCServiceImpl, vpcID string) int {
 	t.Helper()
-	keys, err := svc.rtbKV.Keys()
+	keys, err := svc.rtbKV.Keys(t.Context())
 	if err != nil {
 		return 0
 	}
 	n := 0
 	for _, key := range keys {
-		entry, err := svc.rtbKV.Get(key)
+		entry, err := svc.rtbKV.Get(t.Context(), key)
 		if err != nil {
 			continue
 		}
@@ -655,7 +655,7 @@ func countMainRouteTablesForVPC(t *testing.T, svc *VPCServiceImpl, vpcID string)
 // so subnet and RTB must still land in KV when vpcd is absent.
 func TestEnsureDefaultVPC_NoVpcdResponder(t *testing.T) {
 	_, nc, _ := testutil.StartTestJetStream(t)
-	svc, err := NewVPCServiceImplWithNATS(nil, nc)
+	svc, err := NewVPCServiceImplWithNATS(t.Context(), nil, nc)
 	require.NoError(t, err)
 	// Intentionally NOT calling StubVpcdSGResponder — vpc.create-sg has no
 	// responder, mirroring the bootstrap race.
@@ -676,11 +676,11 @@ func TestEnsureDefaultVPC_NoVpcdResponder(t *testing.T) {
 	assert.Equal(t, info.VpcId, *subDesc.Subnets[0].VpcId)
 
 	require.NotNil(t, svc.rtbKV)
-	rtbKeys, err := svc.rtbKV.Keys()
+	rtbKeys, err := svc.rtbKV.Keys(t.Context())
 	require.NoError(t, err)
 	foundMainRTB := false
 	for _, k := range rtbKeys {
-		entry, err := svc.rtbKV.Get(k)
+		entry, err := svc.rtbKV.Get(t.Context(), k)
 		if err != nil {
 			continue
 		}
@@ -700,7 +700,7 @@ func TestEnsureDefaultVPC_NoVpcdResponder(t *testing.T) {
 
 	// Default SG record is best-effort; KV write happens before the synchronous
 	// vpcd round-trip, so the record should still be present.
-	sgKeys, err := svc.sgKV.Keys()
+	sgKeys, err := svc.sgKV.Keys(t.Context())
 	require.NoError(t, err)
 	assert.NotEmpty(t, sgKeys, "default SG record must persist in KV for vpcd reconciler to converge")
 }
@@ -907,7 +907,7 @@ func TestEnsureDefaultVPC_WithConfigAZ(t *testing.T) {
 	testutil.StubVpcdSGResponder(t, nc)
 
 	cfg := &config.Config{AZ: "us-west-2b"}
-	svc, err := NewVPCServiceImplWithNATS(cfg, nc)
+	svc, err := NewVPCServiceImplWithNATS(t.Context(), cfg, nc)
 	require.NoError(t, err)
 
 	_, err = svc.EnsureDefaultVPC(testAccountID)

spinifex/handlers/ecr/meta.go

@@ -12,8 +12,8 @@ import (
 	"sync"
 	"time"
 
-	"github.com/mulgadc/spinifex/spinifex/utils"
-	"github.com/nats-io/nats.go"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // ErrNotFound is returned by MetaStore reads when the requested record is absent.
@@ -109,26 +109,26 @@ type MetaStore interface {
 // KVMetaStore is the JetStream-KV-backed MetaStore. Per-account buckets are
 // created lazily on first write and cached per process.
 type KVMetaStore struct {
-	js      nats.JetStreamContext
+	js      jetstream.JetStream
 	mu      sync.Mutex
-	buckets map[string]nats.KeyValue
+	buckets map[string]jetstream.KeyValue
 }
 
 var _ MetaStore = (*KVMetaStore)(nil)
 
-// NewKVMetaStore constructs a KVMetaStore over the supplied JetStream context.
-func NewKVMetaStore(js nats.JetStreamContext) *KVMetaStore {
-	return &KVMetaStore{js: js, buckets: make(map[string]nats.KeyValue)}
+// NewKVMetaStore constructs a KVMetaStore over the supplied JetStream handle.
+func NewKVMetaStore(js jetstream.JetStream) *KVMetaStore {
+	return &KVMetaStore{js: js, buckets: make(map[string]jetstream.KeyValue)}
 }
 
-func (s *KVMetaStore) bucket(accountID string) (nats.KeyValue, error) {
+func (s *KVMetaStore) bucket(ctx context.Context, accountID string) (jetstream.KeyValue, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if kv, ok := s.buckets[accountID]; ok {
 		return kv, nil
 	}
 	name := KVAccountBucket(accountID)
-	kv, err := utils.GetOrCreateKVBucket(s.js, name, KVBucketAccountHistory)
+	kv, err := kvutil.GetOrCreateBucket(ctx, s.js, name, KVBucketAccountHistory)
 	if err != nil {
 		return nil, fmt.Errorf("ecr: open account bucket %s: %w", name, err)
 	}
@@ -136,8 +136,8 @@ func (s *KVMetaStore) bucket(accountID string) (nats.KeyValue, error) {
 	return kv, nil
 }
 
-func (s *KVMetaStore) PutRepo(_ context.Context, accountID string, meta RepoMeta) error {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) PutRepo(ctx context.Context, accountID string, meta RepoMeta) error {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return err
 	}
@@ -145,16 +145,16 @@ func (s *KVMetaStore) PutRepo(_ context.Context, accountID string, meta RepoMeta
 	if err != nil {
 		return err
 	}
-	_, err = kv.Put(KVRepoMetaKey(meta.Name), data)
+	_, err = kv.Put(ctx, KVRepoMetaKey(meta.Name), data)
 	return err
 }
 
-func (s *KVMetaStore) GetRepo(_ context.Context, accountID, repo string) (RepoMeta, error) {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) GetRepo(ctx context.Context, accountID, repo string) (RepoMeta, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return RepoMeta{}, err
 	}
-	entry, err := kv.Get(KVRepoMetaKey(repo))
+	entry, err := kv.Get(ctx, KVRepoMetaKey(repo))
 	if err != nil {
 		return RepoMeta{}, mapKVErr(err)
 	}
@@ -165,14 +165,14 @@ func (s *KVMetaStore) GetRepo(_ context.Context, accountID, repo string) (RepoMe
 	return m, nil
 }
 
-func (s *KVMetaStore) ListRepos(_ context.Context, accountID string) ([]string, error) {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) ListRepos(ctx context.Context, accountID string) ([]string, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
-	keys, err := kv.Keys()
+	keys, err := kv.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, nil
 		}
 		return nil, err
@@ -195,17 +195,17 @@ func (s *KVMetaStore) ListRepos(_ context.Context, accountID string) ([]string,
 // tags, and all manifest records. Predastore blob garbage collection is out of
 // scope here (deferred); only the per-account KV records are removed. Returns
 // ErrNotFound when the repository meta is absent.
-func (s *KVMetaStore) DeleteRepo(_ context.Context, accountID, repo string) error {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) DeleteRepo(ctx context.Context, accountID, repo string) error {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return err
 	}
-	if _, err := kv.Get(KVRepoMetaKey(repo)); err != nil {
+	if _, err := kv.Get(ctx, KVRepoMetaKey(repo)); err != nil {
 		return mapKVErr(err)
 	}
-	keys, err := kv.Keys()
+	keys, err := kv.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil
 		}
 		return err
@@ -215,7 +215,7 @@ func (s *KVMetaStore) DeleteRepo(_ context.Context, accountID, repo string) erro
 	for _, k := range keys {
 		if k == metaKey || k == policyKey || k == lifecycleKey ||
 			strings.HasPrefix(k, tagsPrefix) || strings.HasPrefix(k, manifestsPrefix) {
-			if err := kv.Delete(k); err != nil {
+			if err := kv.Delete(ctx, k); err != nil {
 				return mapKVErr(err)
 			}
 		}
@@ -223,14 +223,14 @@ func (s *KVMetaStore) DeleteRepo(_ context.Context, accountID, repo string) erro
 	return nil
 }
 
-func (s *KVMetaStore) ListManifests(_ context.Context, accountID, repo string) ([]string, error) {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) ListManifests(ctx context.Context, accountID, repo string) ([]string, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
-	keys, err := kv.Keys()
+	keys, err := kv.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, nil
 		}
 		return nil, err
@@ -246,127 +246,127 @@ func (s *KVMetaStore) ListManifests(_ context.Context, accountID, repo string) (
 	return digests, nil
 }
 
-func (s *KVMetaStore) PutRepoPolicy(_ context.Context, accountID, repo string, policyText []byte) error {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) PutRepoPolicy(ctx context.Context, accountID, repo string, policyText []byte) error {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return err
 	}
-	_, err = kv.Put(KVRepoPolicyKey(repo), policyText)
+	_, err = kv.Put(ctx, KVRepoPolicyKey(repo), policyText)
 	return err
 }
 
-func (s *KVMetaStore) GetRepoPolicy(_ context.Context, accountID, repo string) ([]byte, error) {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) GetRepoPolicy(ctx context.Context, accountID, repo string) ([]byte, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
-	entry, err := kv.Get(KVRepoPolicyKey(repo))
+	entry, err := kv.Get(ctx, KVRepoPolicyKey(repo))
 	if err != nil {
 		return nil, mapKVErr(err)
 	}
 	return entry.Value(), nil
 }
 
-func (s *KVMetaStore) DeleteRepoPolicy(_ context.Context, accountID, repo string) ([]byte, error) {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) DeleteRepoPolicy(ctx context.Context, accountID, repo string) ([]byte, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 	// JetStream KV Delete is silently idempotent, so the value is read first to
 	// surface a real not-found and return the deleted document to the caller.
-	entry, err := kv.Get(KVRepoPolicyKey(repo))
+	entry, err := kv.Get(ctx, KVRepoPolicyKey(repo))
 	if err != nil {
 		return nil, mapKVErr(err)
 	}
-	if err := kv.Delete(KVRepoPolicyKey(repo)); err != nil {
+	if err := kv.Delete(ctx, KVRepoPolicyKey(repo)); err != nil {
 		return nil, mapKVErr(err)
 	}
 	return entry.Value(), nil
 }
 
-func (s *KVMetaStore) PutLifecyclePolicy(_ context.Context, accountID, repo string, policyText []byte) error {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) PutLifecyclePolicy(ctx context.Context, accountID, repo string, policyText []byte) error {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return err
 	}
-	_, err = kv.Put(KVLifecyclePolicyKey(repo), policyText)
+	_, err = kv.Put(ctx, KVLifecyclePolicyKey(repo), policyText)
 	return err
 }
 
-func (s *KVMetaStore) GetLifecyclePolicy(_ context.Context, accountID, repo string) ([]byte, error) {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) GetLifecyclePolicy(ctx context.Context, accountID, repo string) ([]byte, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
-	entry, err := kv.Get(KVLifecyclePolicyKey(repo))
+	entry, err := kv.Get(ctx, KVLifecyclePolicyKey(repo))
 	if err != nil {
 		return nil, mapKVErr(err)
 	}
 	return entry.Value(), nil
 }
 
-func (s *KVMetaStore) DeleteLifecyclePolicy(_ context.Context, accountID, repo string) ([]byte, error) {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) DeleteLifecyclePolicy(ctx context.Context, accountID, repo string) ([]byte, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 	// JetStream KV Delete is silently idempotent, so the value is read first to
 	// surface a real not-found and return the deleted document to the caller.
-	entry, err := kv.Get(KVLifecyclePolicyKey(repo))
+	entry, err := kv.Get(ctx, KVLifecyclePolicyKey(repo))
 	if err != nil {
 		return nil, mapKVErr(err)
 	}
-	if err := kv.Delete(KVLifecyclePolicyKey(repo)); err != nil {
+	if err := kv.Delete(ctx, KVLifecyclePolicyKey(repo)); err != nil {
 		return nil, mapKVErr(err)
 	}
 	return entry.Value(), nil
 }
 
-func (s *KVMetaStore) PutTag(_ context.Context, accountID, repo, tag, digest string) error {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) PutTag(ctx context.Context, accountID, repo, tag, digest string) error {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return err
 	}
-	_, err = kv.Put(KVTagKey(repo, tag), []byte(digest))
+	_, err = kv.Put(ctx, KVTagKey(repo, tag), []byte(digest))
 	return err
 }
 
-func (s *KVMetaStore) GetTag(_ context.Context, accountID, repo, tag string) (string, error) {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) GetTag(ctx context.Context, accountID, repo, tag string) (string, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return "", err
 	}
-	entry, err := kv.Get(KVTagKey(repo, tag))
+	entry, err := kv.Get(ctx, KVTagKey(repo, tag))
 	if err != nil {
 		return "", mapKVErr(err)
 	}
 	return string(entry.Value()), nil
 }
 
-func (s *KVMetaStore) DeleteTag(_ context.Context, accountID, repo, tag string) error {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) DeleteTag(ctx context.Context, accountID, repo, tag string) error {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return err
 	}
 	// JetStream KV Delete is silently idempotent, so existence is checked first
 	// to surface a real not-found to the caller.
-	if _, err := kv.Get(KVTagKey(repo, tag)); err != nil {
+	if _, err := kv.Get(ctx, KVTagKey(repo, tag)); err != nil {
 		return mapKVErr(err)
 	}
-	if err := kv.Delete(KVTagKey(repo, tag)); err != nil {
+	if err := kv.Delete(ctx, KVTagKey(repo, tag)); err != nil {
 		return mapKVErr(err)
 	}
 	return nil
 }
 
-func (s *KVMetaStore) ListTags(_ context.Context, accountID, repo string) ([]string, error) {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) ListTags(ctx context.Context, accountID, repo string) ([]string, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
-	keys, err := kv.Keys()
+	keys, err := kv.Keys(ctx)
 	if err != nil {
-		if errors.Is(err, nats.ErrNoKeysFound) {
+		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return nil, nil
 		}
 		return nil, err
@@ -382,8 +382,8 @@ func (s *KVMetaStore) ListTags(_ context.Context, accountID, repo string) ([]str
 	return tags, nil
 }
 
-func (s *KVMetaStore) PutManifestMeta(_ context.Context, accountID, repo string, meta ManifestMeta) error {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) PutManifestMeta(ctx context.Context, accountID, repo string, meta ManifestMeta) error {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return err
 	}
@@ -391,16 +391,16 @@ func (s *KVMetaStore) PutManifestMeta(_ context.Context, accountID, repo string,
 	if err != nil {
 		return err
 	}
-	_, err = kv.Put(KVManifestKey(repo, meta.Digest), data)
+	_, err = kv.Put(ctx, KVManifestKey(repo, meta.Digest), data)
 	return err
 }
 
-func (s *KVMetaStore) GetManifestMeta(_ context.Context, accountID, repo, digest string) (ManifestMeta, error) {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) GetManifestMeta(ctx context.Context, accountID, repo, digest string) (ManifestMeta, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return ManifestMeta{}, err
 	}
-	entry, err := kv.Get(KVManifestKey(repo, digest))
+	entry, err := kv.Get(ctx, KVManifestKey(repo, digest))
 	if err != nil {
 		return ManifestMeta{}, mapKVErr(err)
 	}
@@ -411,19 +411,19 @@ func (s *KVMetaStore) GetManifestMeta(_ context.Context, accountID, repo, digest
 	return m, nil
 }
 
-func (s *KVMetaStore) DeleteManifestMeta(_ context.Context, accountID, repo, digest string) error {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) DeleteManifestMeta(ctx context.Context, accountID, repo, digest string) error {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return err
 	}
-	if _, err := kv.Get(KVManifestKey(repo, digest)); err != nil {
+	if _, err := kv.Get(ctx, KVManifestKey(repo, digest)); err != nil {
 		return mapKVErr(err)
 	}
-	return kv.Delete(KVManifestKey(repo, digest))
+	return kv.Delete(ctx, KVManifestKey(repo, digest))
 }
 
-func (s *KVMetaStore) PutUpload(_ context.Context, accountID, uploadID string, state UploadState) (uint64, error) {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) PutUpload(ctx context.Context, accountID, uploadID string, state UploadState) (uint64, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return 0, err
 	}
@@ -431,15 +431,15 @@ func (s *KVMetaStore) PutUpload(_ context.Context, accountID, uploadID string, s
 	if err != nil {
 		return 0, err
 	}
-	return kv.Put(KVUploadKey(uploadID), data)
+	return kv.Put(ctx, KVUploadKey(uploadID), data)
 }
 
-func (s *KVMetaStore) GetUpload(_ context.Context, accountID, uploadID string) (UploadState, uint64, error) {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) GetUpload(ctx context.Context, accountID, uploadID string) (UploadState, uint64, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return UploadState{}, 0, err
 	}
-	entry, err := kv.Get(KVUploadKey(uploadID))
+	entry, err := kv.Get(ctx, KVUploadKey(uploadID))
 	if err != nil {
 		return UploadState{}, 0, mapKVErr(err)
 	}
@@ -450,8 +450,8 @@ func (s *KVMetaStore) GetUpload(_ context.Context, accountID, uploadID string) (
 	return st, entry.Revision(), nil
 }
 
-func (s *KVMetaStore) UpdateUpload(_ context.Context, accountID, uploadID string, state UploadState, rev uint64) (uint64, error) {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) UpdateUpload(ctx context.Context, accountID, uploadID string, state UploadState, rev uint64) (uint64, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return 0, err
 	}
@@ -459,9 +459,9 @@ func (s *KVMetaStore) UpdateUpload(_ context.Context, accountID, uploadID string
 	if err != nil {
 		return 0, err
 	}
-	newRev, err := kv.Update(KVUploadKey(uploadID), data, rev)
+	newRev, err := kv.Update(ctx, KVUploadKey(uploadID), data, rev)
 	if err != nil {
-		if errors.Is(err, nats.ErrKeyExists) {
+		if errors.Is(err, jetstream.ErrKeyExists) {
 			return 0, ErrConflict
 		}
 		return 0, err
@@ -469,17 +469,17 @@ func (s *KVMetaStore) UpdateUpload(_ context.Context, accountID, uploadID string
 	return newRev, nil
 }
 
-func (s *KVMetaStore) DeleteUpload(_ context.Context, accountID, uploadID string) error {
-	kv, err := s.bucket(accountID)
+func (s *KVMetaStore) DeleteUpload(ctx context.Context, accountID, uploadID string) error {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return err
 	}
 	// JetStream KV Delete is silently idempotent, so existence is checked first
 	// to surface a real not-found to the caller.
-	if _, err := kv.Get(KVUploadKey(uploadID)); err != nil {
+	if _, err := kv.Get(ctx, KVUploadKey(uploadID)); err != nil {
 		return mapKVErr(err)
 	}
-	if err := kv.Delete(KVUploadKey(uploadID)); err != nil {
+	if err := kv.Delete(ctx, KVUploadKey(uploadID)); err != nil {
 		return mapKVErr(err)
 	}
 	return nil
@@ -763,7 +763,7 @@ func (m *MemoryMetaStore) DeleteUpload(_ context.Context, accountID, uploadID st
 
 // mapKVErr normalizes JetStream KV errors to the MetaStore error vocabulary.
 func mapKVErr(err error) error {
-	if errors.Is(err, nats.ErrKeyNotFound) {
+	if errors.Is(err, jetstream.ErrKeyNotFound) {
 		return ErrNotFound
 	}
 	return err

spinifex/handlers/ecr/meta_kv_test.go

@@ -14,7 +14,8 @@ import (
 // TestKVMetaStore_FullCycle exercises the JetStream-backed MetaStore against an
 // embedded JetStream server, covering the production CAS path.
 func TestKVMetaStore_FullCycle(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 	store := NewKVMetaStore(js)
 	const acct = "000000000000"
 
@@ -90,7 +91,8 @@ func TestKVMetaStore_FullCycle(t *testing.T) {
 // tags, and manifests, leaves a same-prefix nested repo intact, and reports
 // ErrNotFound for an absent repo.
 func TestKVMetaStore_DeleteRepo(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 	store := NewKVMetaStore(js)
 	const acct = "000000000000"
 	const dig = "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
@@ -133,7 +135,8 @@ func TestKVMetaStore_DeleteRepo(t *testing.T) {
 // JetStream-backed store and proves the policy key shares the repo bucket
 // without being mistaken for a repository by ListRepos.
 func TestKVMetaStore_RepoPolicy(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 	store := NewKVMetaStore(js)
 	const acct = "000000000000"
 	const policy = `{"Version":"2012-10-17","Statement":[]}`
@@ -164,7 +167,8 @@ func TestKVMetaStore_RepoPolicy(t *testing.T) {
 // against the JetStream-backed store and proves the lifecycle key shares the
 // repo bucket without being mistaken for a repository by ListRepos.
 func TestKVMetaStore_LifecyclePolicy(t *testing.T) {
-	_, _, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 	store := NewKVMetaStore(js)
 	const acct = "000000000000"
 	const policy = `{"rules":[{"rulePriority":1,"selection":{"tagStatus":"untagged","countType":"sinceImagePushed","countUnit":"days","countNumber":14},"action":{"type":"expire"}}]}`

spinifex/handlers/ecr/service_impl.go

@@ -4,7 +4,7 @@ import (
 	"context"
 	"errors"
 
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // MetaServiceImpl is the daemon-side MetaService. It owns the backing MetaStore
@@ -23,7 +23,7 @@ func NewMetaServiceImpl(store MetaStore) *MetaServiceImpl {
 }
 
 // NewKVMetaService builds a MetaService backed by per-account JetStream KV.
-func NewKVMetaService(js nats.JetStreamContext) *MetaServiceImpl {
+func NewKVMetaService(js jetstream.JetStream) *MetaServiceImpl {
 	return &MetaServiceImpl{store: NewKVMetaStore(js)}
 }
 

spinifex/handlers/ecr/service_test.go

@@ -56,7 +56,8 @@ func TestMetaServiceImpl_FlagMapping(t *testing.T) {
 // TestNATSMetaStore_RepoPolicy drives the policy passthrough round-trip
 // (put/get/delete + not-found mapping) end to end over the embedded server.
 func TestNATSMetaStore_RepoPolicy(t *testing.T) {
-	_, nc, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 	svc := NewKVMetaService(js)
 
 	serveMeta(t, nc, SubjectPolicyPut, svc.PolicyPut)
@@ -88,7 +89,8 @@ func TestNATSMetaStore_RepoPolicy(t *testing.T) {
 // TestNATSMetaStore_LifecyclePolicy drives the lifecycle-policy passthrough
 // round-trip (put/get/delete + not-found mapping) end to end.
 func TestNATSMetaStore_LifecyclePolicy(t *testing.T) {
-	_, nc, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 	svc := NewKVMetaService(js)
 
 	serveMeta(t, nc, SubjectLifecyclePut, svc.LifecyclePut)
@@ -143,7 +145,8 @@ func serveMeta[I any, O any](t *testing.T, nc *nats.Conn, subject string, fn fun
 // a KV-backed daemon service over an embedded NATS/JetStream server, proving the
 // full request/reply path including ErrNotFound/ErrConflict reconstruction.
 func TestNATSMetaStore_RoundTrip(t *testing.T) {
-	_, nc, js := testutil.StartTestJetStream(t)
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
 	svc := NewKVMetaService(js)
 
 	serveMeta(t, nc, SubjectRepoCreate, svc.RepoCreate)

spinifex/handlers/ecs/capacity_providers.go

@@ -13,17 +13,17 @@ import (
 // PutClusterCapacityProviders persists the cluster's capacity-provider names
 // and default strategy. Accepted and stored only: no scheduler coupling, no
 // scale loop (a separate follow-on binds this to an ASG primitive).
-func (s *Service) PutClusterCapacityProviders(_ context.Context, input *ecs.PutClusterCapacityProvidersInput, accountID string) (*ecs.PutClusterCapacityProvidersOutput, error) {
+func (s *Service) PutClusterCapacityProviders(ctx context.Context, input *ecs.PutClusterCapacityProvidersInput, accountID string) (*ecs.PutClusterCapacityProvidersOutput, error) {
 	if input == nil {
 		return nil, errors.New(awserrors.ErrorECSInvalidParameter)
 	}
 	cluster := clusterShortName(aws.StringValue(input.Cluster))
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 	var rec ClusterRecord
-	found, err := getJSON(kv, ClusterMetaKey(cluster), &rec)
+	found, err := getJSON(ctx, kv, ClusterMetaKey(cluster), &rec)
 	if err != nil {
 		return nil, err
 	}
@@ -33,7 +33,7 @@ func (s *Service) PutClusterCapacityProviders(_ context.Context, input *ecs.PutC
 
 	rec.CapacityProviders = awsStringSlice(input.CapacityProviders)
 	rec.DefaultCapacityProviderStrategy = strategyFromAWS(input.DefaultCapacityProviderStrategy)
-	if err := putJSON(kv, ClusterMetaKey(cluster), &rec); err != nil {
+	if err := putJSON(ctx, kv, ClusterMetaKey(cluster), &rec); err != nil {
 		return nil, err
 	}
 	return &ecs.PutClusterCapacityProvidersOutput{Cluster: rec.toAWS()}, nil
@@ -41,19 +41,19 @@ func (s *Service) PutClusterCapacityProviders(_ context.Context, input *ecs.PutC
 
 // CreateCapacityProvider persists a new capacity provider. Idempotent:
 // re-creating an existing name returns the existing record.
-func (s *Service) CreateCapacityProvider(_ context.Context, input *ecs.CreateCapacityProviderInput, accountID string) (*ecs.CreateCapacityProviderOutput, error) {
+func (s *Service) CreateCapacityProvider(ctx context.Context, input *ecs.CreateCapacityProviderInput, accountID string) (*ecs.CreateCapacityProviderOutput, error) {
 	if input == nil || aws.StringValue(input.Name) == "" || input.AutoScalingGroupProvider == nil ||
 		aws.StringValue(input.AutoScalingGroupProvider.AutoScalingGroupArn) == "" {
 		return nil, errors.New(awserrors.ErrorECSInvalidParameter)
 	}
 	name := aws.StringValue(input.Name)
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
 	var existing CapacityProviderRecord
-	found, err := getJSON(kv, CapacityProviderKey(name), &existing)
+	found, err := getJSON(ctx, kv, CapacityProviderKey(name), &existing)
 	if err != nil {
 		return nil, err
 	}
@@ -69,7 +69,7 @@ func (s *Service) CreateCapacityProvider(_ context.Context, input *ecs.CreateCap
 		Tags:                     tagsToMap(input.Tags),
 		CreatedAt:                time.Now().UTC(),
 	}
-	if err := putJSON(kv, CapacityProviderKey(name), &rec); err != nil {
+	if err := putJSON(ctx, kv, CapacityProviderKey(name), &rec); err != nil {
 		return nil, err
 	}
 	return &ecs.CreateCapacityProviderOutput{CapacityProvider: rec.toAWS()}, nil
@@ -77,8 +77,8 @@ func (s *Service) CreateCapacityProvider(_ context.Context, input *ecs.CreateCap
 
 // DescribeCapacityProviders returns the named providers, or every provider in
 // the account when none are named. Unknown names surface as failures.
-func (s *Service) DescribeCapacityProviders(_ context.Context, input *ecs.DescribeCapacityProvidersInput, accountID string) (*ecs.DescribeCapacityProvidersOutput, error) {
-	kv, err := s.bucket(accountID)
+func (s *Service) DescribeCapacityProviders(ctx context.Context, input *ecs.DescribeCapacityProvidersInput, accountID string) (*ecs.DescribeCapacityProvidersOutput, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -86,13 +86,13 @@ func (s *Service) DescribeCapacityProviders(_ context.Context, input *ecs.Descri
 
 	names := awsStringSlice(input.CapacityProviders)
 	if len(names) == 0 {
-		keys, kerr := keysWithPrefix(kv, CapacityProvidersPrefix())
+		keys, kerr := keysWithPrefix(ctx, kv, CapacityProvidersPrefix())
 		if kerr != nil {
 			return nil, kerr
 		}
 		for _, k := range keys {
 			var rec CapacityProviderRecord
-			ok, gerr := getJSON(kv, k, &rec)
+			ok, gerr := getJSON(ctx, kv, k, &rec)
 			if gerr != nil {
 				return nil, gerr
 			}
@@ -106,7 +106,7 @@ func (s *Service) DescribeCapacityProviders(_ context.Context, input *ecs.Descri
 	for _, ref := range names {
 		name := capacityProviderShortName(ref)
 		var rec CapacityProviderRecord
-		ok, gerr := getJSON(kv, CapacityProviderKey(name), &rec)
+		ok, gerr := getJSON(ctx, kv, CapacityProviderKey(name), &rec)
 		if gerr != nil {
 			return nil, gerr
 		}
@@ -122,24 +122,24 @@ func (s *Service) DescribeCapacityProviders(_ context.Context, input *ecs.Descri
 // DeleteCapacityProvider removes a capacity provider record. Clusters that
 // still reference the name by CapacityProviders are left untouched (v1 does
 // not enforce referential integrity here).
-func (s *Service) DeleteCapacityProvider(_ context.Context, input *ecs.DeleteCapacityProviderInput, accountID string) (*ecs.DeleteCapacityProviderOutput, error) {
+func (s *Service) DeleteCapacityProvider(ctx context.Context, input *ecs.DeleteCapacityProviderInput, accountID string) (*ecs.DeleteCapacityProviderOutput, error) {
 	if input == nil || aws.StringValue(input.CapacityProvider) == "" {
 		return nil, errors.New(awserrors.ErrorECSInvalidParameter)
 	}
 	name := capacityProviderShortName(aws.StringValue(input.CapacityProvider))
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 	var rec CapacityProviderRecord
-	found, err := getJSON(kv, CapacityProviderKey(name), &rec)
+	found, err := getJSON(ctx, kv, CapacityProviderKey(name), &rec)
 	if err != nil {
 		return nil, err
 	}
 	if !found {
 		return nil, errors.New(awserrors.ErrorECSInvalidParameter)
 	}
-	if err := kv.Delete(CapacityProviderKey(name)); err != nil {
+	if err := kv.Delete(ctx, CapacityProviderKey(name)); err != nil {
 		return nil, err
 	}
 	return &ecs.DeleteCapacityProviderOutput{CapacityProvider: rec.toAWS()}, nil

spinifex/handlers/ecs/cluster_teardown.go

@@ -7,7 +7,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ecs"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // clusterKeyPrefix is the KV subtree holding all of a cluster's records.
@@ -23,12 +23,12 @@ func clusterKeyPrefix(cluster string) string {
 // the cluster with Status INACTIVE.
 func (s *Service) DeleteCluster(ctx context.Context, input *ecs.DeleteClusterInput, accountID string) (*ecs.DeleteClusterOutput, error) {
 	cluster := clusterShortName(aws.StringValue(input.Cluster))
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 	var rec ClusterRecord
-	found, err := getJSON(kv, ClusterMetaKey(cluster), &rec)
+	found, err := getJSON(ctx, kv, ClusterMetaKey(cluster), &rec)
 	if err != nil {
 		return nil, err
 	}
@@ -36,7 +36,7 @@ func (s *Service) DeleteCluster(ctx context.Context, input *ecs.DeleteClusterInp
 		return nil, errors.New(awserrors.ErrorECSClusterNotFound)
 	}
 
-	tasks, err := s.listTaskRecords(kv, cluster)
+	tasks, err := s.listTaskRecords(ctx, kv, cluster)
 	if err != nil {
 		return nil, err
 	}
@@ -44,7 +44,7 @@ func (s *Service) DeleteCluster(ctx context.Context, input *ecs.DeleteClusterInp
 		s.forceStopTask(ctx, kv, accountID, &tasks[i], "Cluster deleted")
 	}
 
-	if err := deleteKeysWithPrefix(kv, clusterKeyPrefix(cluster)); err != nil {
+	if err := deleteKeysWithPrefix(ctx, kv, clusterKeyPrefix(cluster)); err != nil {
 		return nil, err
 	}
 
@@ -60,12 +60,12 @@ func (s *Service) DeleteCluster(ctx context.Context, input *ecs.DeleteClusterInp
 func (s *Service) DeregisterContainerInstance(ctx context.Context, input *ecs.DeregisterContainerInstanceInput, accountID string) (*ecs.DeregisterContainerInstanceOutput, error) {
 	cluster := clusterShortName(aws.StringValue(input.Cluster))
 	id := containerInstanceShortID(aws.StringValue(input.ContainerInstance))
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 	var rec InstanceRecord
-	found, err := getJSON(kv, InstanceKey(cluster, id), &rec)
+	found, err := getJSON(ctx, kv, InstanceKey(cluster, id), &rec)
 	if err != nil {
 		return nil, err
 	}
@@ -73,7 +73,7 @@ func (s *Service) DeregisterContainerInstance(ctx context.Context, input *ecs.De
 		return nil, errors.New(awserrors.ErrorECSInvalidParameter)
 	}
 
-	active, err := s.instanceActiveTasks(kv, cluster, id)
+	active, err := s.instanceActiveTasks(ctx, kv, cluster, id)
 	if err != nil {
 		return nil, err
 	}
@@ -84,10 +84,10 @@ func (s *Service) DeregisterContainerInstance(ctx context.Context, input *ecs.De
 		s.forceStopTask(ctx, kv, accountID, &active[i], "Container instance deregistered")
 	}
 
-	if derr := deleteKeysWithPrefix(kv, AssignmentsPrefix(cluster, id)); derr != nil {
+	if derr := deleteKeysWithPrefix(ctx, kv, AssignmentsPrefix(cluster, id)); derr != nil {
 		return nil, derr
 	}
-	if derr := kv.Delete(InstanceKey(cluster, id)); derr != nil {
+	if derr := kv.Delete(ctx, InstanceKey(cluster, id)); derr != nil {
 		return nil, derr
 	}
 	rec.Status = ClusterStatusInactive
@@ -104,7 +104,7 @@ func (s *Service) UpdateContainerInstancesState(ctx context.Context, input *ecs.
 	if status != InstanceStatusActive && status != InstanceStatusDraining {
 		return nil, errors.New(awserrors.ErrorECSInvalidParameter)
 	}
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -112,7 +112,7 @@ func (s *Service) UpdateContainerInstancesState(ctx context.Context, input *ecs.
 	for _, ref := range awsStringSlice(input.ContainerInstances) {
 		id := containerInstanceShortID(ref)
 		var rec InstanceRecord
-		found, gerr := getJSON(kv, InstanceKey(cluster, id), &rec)
+		found, gerr := getJSON(ctx, kv, InstanceKey(cluster, id), &rec)
 		if gerr != nil {
 			return nil, gerr
 		}
@@ -122,7 +122,7 @@ func (s *Service) UpdateContainerInstancesState(ctx context.Context, input *ecs.
 		}
 		rec.Status = status
 		rec.Reaped = false
-		if perr := putJSON(kv, InstanceKey(cluster, id), &rec); perr != nil {
+		if perr := putJSON(ctx, kv, InstanceKey(cluster, id), &rec); perr != nil {
 			return nil, perr
 		}
 		if status == InstanceStatusDraining {
@@ -134,8 +134,8 @@ func (s *Service) UpdateContainerInstancesState(ctx context.Context, input *ecs.
 }
 
 // instanceActiveTasks returns a cluster's non-STOPPED tasks placed on instanceID.
-func (s *Service) instanceActiveTasks(kv nats.KeyValue, cluster, instanceID string) ([]TaskRecord, error) {
-	all, err := s.listTaskRecords(kv, cluster)
+func (s *Service) instanceActiveTasks(ctx context.Context, kv jetstream.KeyValue, cluster, instanceID string) ([]TaskRecord, error) {
+	all, err := s.listTaskRecords(ctx, kv, cluster)
 	if err != nil {
 		return nil, err
 	}
@@ -151,8 +151,8 @@ func (s *Service) instanceActiveTasks(kv nats.KeyValue, cluster, instanceID stri
 // drainInstanceServiceTasks force-stops the instance's service-owned tasks on an
 // intentional DRAINING; the service reconciler then relaunches them on another
 // instance. Standalone tasks are left running (AWS DRAINING semantics).
-func (s *Service) drainInstanceServiceTasks(ctx context.Context, kv nats.KeyValue, accountID, cluster, instanceID string) {
-	active, err := s.instanceActiveTasks(kv, cluster, instanceID)
+func (s *Service) drainInstanceServiceTasks(ctx context.Context, kv jetstream.KeyValue, accountID, cluster, instanceID string) {
+	active, err := s.instanceActiveTasks(ctx, kv, cluster, instanceID)
 	if err != nil {
 		return
 	}

spinifex/handlers/ecs/cluster_teardown_test.go

@@ -26,7 +26,7 @@ func TestService_DeleteCluster_CascadesAndSweeps(t *testing.T) {
 	list, err := svc.ListClusters(context.Background(), &ecs.ListClustersInput{}, testAccountID)
 	require.NoError(t, err)
 	assert.Empty(t, list.ClusterArns)
-	keys, err := keysWithPrefix(kv, clusterKeyPrefix("web"))
+	keys, err := keysWithPrefix(t.Context(), kv, clusterKeyPrefix("web"))
 	require.NoError(t, err)
 	assert.Empty(t, keys)
 }
@@ -81,7 +81,7 @@ func TestService_DeregisterContainerInstance_ForceStopsAndRemoves(t *testing.T)
 	require.NoError(t, err)
 	assert.Equal(t, ClusterStatusInactive, aws.StringValue(out.ContainerInstance.Status))
 
-	active, err := svc.instanceActiveTasks(kv, "web", "i-1")
+	active, err := svc.instanceActiveTasks(t.Context(), kv, "web", "i-1")
 	require.NoError(t, err)
 	assert.Empty(t, active)
 
@@ -111,7 +111,7 @@ func TestService_UpdateContainerInstancesState_DrainsServiceTasks(t *testing.T)
 	assert.Equal(t, InstanceStatusDraining, aws.StringValue(out.ContainerInstances[0].Status))
 
 	// Service-owned tasks on the drained instance are stopped.
-	live, err := svc.listServiceTasks(kv, "web", "web")
+	live, err := svc.listServiceTasks(t.Context(), kv, "web", "web")
 	require.NoError(t, err)
 	assert.Empty(t, live)
 }

spinifex/handlers/ecs/deployments_test.go

@@ -7,26 +7,26 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ecs"
 	"github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 // reloadService re-reads a service record from KV so a test sees writes made by
 // the reconciler / failure accounting.
-func reloadService(t *testing.T, kv nats.KeyValue, cluster, name string) *ServiceRecord {
+func reloadService(t *testing.T, kv jetstream.KeyValue, cluster, name string) *ServiceRecord {
 	t.Helper()
 	var rec ServiceRecord
-	found, err := getJSON(kv, ServiceKey(cluster, name), &rec)
+	found, err := getJSON(t.Context(), kv, ServiceKey(cluster, name), &rec)
 	require.NoError(t, err)
 	require.True(t, found)
 	return &rec
 }
 
 // driveRunning reports RUNNING for every PENDING task of a service.
-func driveRunning(t *testing.T, svc *Service, kv nats.KeyValue, cluster, name string) {
+func driveRunning(t *testing.T, svc *Service, kv jetstream.KeyValue, cluster, name string) {
 	t.Helper()
-	tasks, err := svc.listServiceTasks(kv, cluster, name)
+	tasks, err := svc.listServiceTasks(t.Context(), kv, cluster, name)
 	require.NoError(t, err)
 	for i := range tasks {
 		if tasks[i].LastStatus != TaskStatusPending {
@@ -41,9 +41,9 @@ func driveRunning(t *testing.T, svc *Service, kv nats.KeyValue, cluster, name st
 
 // failPending reports STOPPED (never-RUNNING) for every PENDING task, simulating
 // tasks that fail to start — the deployment circuit-breaker signal.
-func failPending(t *testing.T, svc *Service, kv nats.KeyValue, cluster, name string) {
+func failPending(t *testing.T, svc *Service, kv jetstream.KeyValue, cluster, name string) {
 	t.Helper()
-	tasks, err := svc.listServiceTasks(kv, cluster, name)
+	tasks, err := svc.listServiceTasks(t.Context(), kv, cluster, name)
 	require.NoError(t, err)
 	for i := range tasks {
 		if tasks[i].LastStatus != TaskStatusPending {
@@ -73,7 +73,7 @@ func TestDeployment_CreateService_SeedsPrimary(t *testing.T) {
 	assert.Equal(t, defaultMinimumHealthyPercent, rec.MinimumHealthyPercent)
 	assert.Equal(t, defaultMaximumPercent, rec.MaximumPercent)
 	// Every launched task carries the primary deployment id in StartedBy.
-	tasks, err := svc.listServiceTasks(kv, "web", "web")
+	tasks, err := svc.listServiceTasks(t.Context(), kv, "web", "web")
 	require.NoError(t, err)
 	require.Len(t, tasks, 2)
 	for i := range tasks {
@@ -110,7 +110,7 @@ func TestDeployment_RollingUpdate_ReplacesOldWithNew(t *testing.T) {
 	assert.NotEqual(t, firstDeployID, rec.DeploymentID)
 	// maximumPercent=200 lets 2 new tasks launch alongside the 2 old running ones.
 	newTasks := 0
-	tasks, err := svc.listServiceTasks(kv, "web", "web")
+	tasks, err := svc.listServiceTasks(t.Context(), kv, "web", "web")
 	require.NoError(t, err)
 	for i := range tasks {
 		if deploymentIDFromStartedBy(tasks[i].StartedBy) == rec.DeploymentID {
@@ -128,7 +128,7 @@ func TestDeployment_RollingUpdate_ReplacesOldWithNew(t *testing.T) {
 	assert.Equal(t, rec.DeploymentID, rec.primaryDeployment().ID)
 
 	// Only new-deployment tasks remain live.
-	tasks, err = svc.listServiceTasks(kv, "web", "web")
+	tasks, err = svc.listServiceTasks(t.Context(), kv, "web", "web")
 	require.NoError(t, err)
 	require.Len(t, tasks, 2)
 	for i := range tasks {
@@ -159,7 +159,7 @@ func TestDeployment_MinimumHealthyPercent_GatesDrain(t *testing.T) {
 
 	// Ceiling of 3 tasks total: 2 old running + 1 new pending. No old task drained
 	// yet because healthy running (2) must not drop below minCount (2).
-	tasks, err := svc.listServiceTasks(kv, "web", "web")
+	tasks, err := svc.listServiceTasks(t.Context(), kv, "web", "web")
 	require.NoError(t, err)
 	assert.Len(t, tasks, 3)
 	rec := reloadService(t, kv, "web", "web")
@@ -215,7 +215,7 @@ func TestDeployment_LegacyServiceSynthesizesPrimary(t *testing.T) {
 		DesiredCount: 1, Status: ServiceStatusActive,
 		SchedulingStrategy: SchedulingStrategyReplica, DeploymentID: "legacy-dep",
 	}
-	require.NoError(t, putJSON(kv, ServiceKey("web", "legacy"), rec))
+	require.NoError(t, putJSON(t.Context(), kv, ServiceKey("web", "legacy"), rec))
 
 	require.NoError(t, svc.reconcileService(context.Background(), kv, testAccountID, rec))
 	reloaded := reloadService(t, kv, "web", "legacy")

spinifex/handlers/ecs/eip_test.go

@@ -37,14 +37,14 @@ func TestAssignTaskPublicIP_EnabledAssociatesAndPersists(t *testing.T) {
 	eips := &stubEIPs{publicIP: "192.0.2.50", allocID: "eipalloc-1"}
 	svc.eips = eips
 
-	require.NoError(t, putJSON(kv, ServiceKey("web", "web"), &ServiceRecord{
+	require.NoError(t, putJSON(t.Context(), kv, ServiceKey("web", "web"), &ServiceRecord{
 		Name: "web", Cluster: "web", AssignPublicIP: "ENABLED",
 	}))
 	task := &TaskRecord{
 		TaskID: "t-1", Cluster: "web", Group: serviceTaskGroup("web"),
 		ENIID: "eni-1", ENIPrivateIP: "172.31.0.8",
 	}
-	require.NoError(t, putJSON(kv, TaskKey("web", "t-1"), task))
+	require.NoError(t, putJSON(t.Context(), kv, TaskKey("web", "t-1"), task))
 
 	svc.assignTaskPublicIP(context.Background(), kv, testAccountID, task)
 
@@ -53,7 +53,7 @@ func TestAssignTaskPublicIP_EnabledAssociatesAndPersists(t *testing.T) {
 	assert.Equal(t, "eipalloc-1", task.ENIEIPAllocationID)
 
 	var persisted TaskRecord
-	found, err := getJSON(kv, TaskKey("web", "t-1"), &persisted)
+	found, err := getJSON(t.Context(), kv, TaskKey("web", "t-1"), &persisted)
 	require.NoError(t, err)
 	require.True(t, found)
 	assert.Equal(t, "192.0.2.50", persisted.ENIPublicIP)
@@ -64,7 +64,7 @@ func TestAssignTaskPublicIP_DisabledNoOp(t *testing.T) {
 	eips := &stubEIPs{publicIP: "192.0.2.50", allocID: "eipalloc-1"}
 	svc.eips = eips
 
-	require.NoError(t, putJSON(kv, ServiceKey("web", "web"), &ServiceRecord{
+	require.NoError(t, putJSON(t.Context(), kv, ServiceKey("web", "web"), &ServiceRecord{
 		Name: "web", Cluster: "web", // AssignPublicIP unset
 	}))
 	task := &TaskRecord{

spinifex/handlers/ecs/eni_test.go

@@ -208,13 +208,13 @@ func TestReaper_Awsvpc_ReleasesENI(t *testing.T) {
 	_, err = svc.RunTask(context.Background(), awsvpcRunInput(), testAccountID)
 	require.NoError(t, err)
 
-	kv, err := svc.bucket(testAccountID)
+	kv, err := svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
 	var rec InstanceRecord
-	_, err = getJSON(kv, InstanceKey("web", "i-1"), &rec)
+	_, err = getJSON(t.Context(), kv, InstanceKey("web", "i-1"), &rec)
 	require.NoError(t, err)
 	rec.LastSeen = time.Now().UTC().Add(-2 * heartbeatTimeout)
-	require.NoError(t, putJSON(kv, InstanceKey("web", "i-1"), &rec))
+	require.NoError(t, putJSON(t.Context(), kv, InstanceKey("web", "i-1"), &rec))
 
 	sc := NewScheduler(nc, svc, "test-holder")
 	sc.reapBucket(context.Background(), kv, testAccountID, time.Now().UTC())

spinifex/handlers/ecs/enumeration_failure_test.go

@@ -0,0 +1,42 @@
+package handlers_ecs
+
+import (
+	"context"
+	"testing"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/ecs"
+	"github.com/stretchr/testify/require"
+)
+
+// TestSchedulerPasses_EnumerationFailureIsReported pins the contract every
+// leader-only sweep depends on: the KV bucket lister closes its channel the same
+// way whether the listing completed or failed, so a pass that ignored the
+// terminal error would read an unreachable JetStream as an empty fleet and
+// report a clean tick while every account went unswept — instances never
+// drained, services never relaunched, stopped tasks never pruned.
+func TestSchedulerPasses_EnumerationFailureIsReported(t *testing.T) {
+	svc, nc := newTestService(t)
+	sc := NewScheduler(nc, svc, "test-holder")
+	_, err := svc.CreateCluster(context.Background(), &ecs.CreateClusterInput{ClusterName: aws.String("web")}, testAccountID)
+	require.NoError(t, err)
+
+	passes := map[string]func(context.Context) error{
+		"reap":                 sc.reap,
+		"sweepStoppedTasks":    sc.sweepStoppedTasks,
+		"reconcileAllServices": svc.reconcileAllServices,
+	}
+	for name, pass := range passes {
+		t.Run(name+"/reachable", func(t *testing.T) {
+			require.NoError(t, pass(t.Context()), "a complete enumeration must complete the pass")
+		})
+	}
+
+	// Closing the connection fails the stream-names request behind the listing.
+	nc.Close()
+	for name, pass := range passes {
+		t.Run(name+"/unreachable", func(t *testing.T) {
+			require.Error(t, pass(t.Context()), "a failed bucket enumeration must not report a completed pass")
+		})
+	}
+}

spinifex/handlers/ecs/gpu_report.go

@@ -28,15 +28,15 @@ type ReportTaskGPUOutput struct {
 // container records by name, so DescribeTasks can surface gpuIds. A missing
 // task or container is a silent no-op — the state-report path already owns
 // the task's lifecycle; this only enriches it.
-func (s *Service) ReportTaskGPU(_ context.Context, input *ReportTaskGPUInput, accountID string) (*ReportTaskGPUOutput, error) {
+func (s *Service) ReportTaskGPU(ctx context.Context, input *ReportTaskGPUInput, accountID string) (*ReportTaskGPUOutput, error) {
 	cluster := clusterShortName(input.Cluster)
 	taskID := taskShortID(input.Task)
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 	var task TaskRecord
-	found, err := getJSON(kv, TaskKey(cluster, taskID), &task)
+	found, err := getJSON(ctx, kv, TaskKey(cluster, taskID), &task)
 	if err != nil {
 		return nil, err
 	}
@@ -55,7 +55,7 @@ func (s *Service) ReportTaskGPU(_ context.Context, input *ReportTaskGPUInput, ac
 		}
 	}
 	if changed {
-		if err := putJSON(kv, TaskKey(cluster, taskID), &task); err != nil {
+		if err := putJSON(ctx, kv, TaskKey(cluster, taskID), &task); err != nil {
 			return nil, err
 		}
 	}

spinifex/handlers/ecs/gpu_report_test.go

@@ -15,9 +15,9 @@ import (
 // container(s) of an existing task record, by container name.
 func TestReportTaskGPU_MergesOntoTaskContainers(t *testing.T) {
 	svc, _ := newTestService(t)
-	kv, err := svc.bucket(testAccountID)
+	kv, err := svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
-	require.NoError(t, putJSON(kv, TaskKey("web", "t-1"), &TaskRecord{
+	require.NoError(t, putJSON(t.Context(), kv, TaskKey("web", "t-1"), &TaskRecord{
 		TaskID: "t-1", Cluster: "web",
 		Containers: []ContainerState{
 			{Name: "web", Status: "RUNNING"},
@@ -33,7 +33,7 @@ func TestReportTaskGPU_MergesOntoTaskContainers(t *testing.T) {
 	assert.Equal(t, "OK", out.Acknowledgment)
 
 	var rec TaskRecord
-	found, err := getJSON(kv, TaskKey("web", "t-1"), &rec)
+	found, err := getJSON(t.Context(), kv, TaskKey("web", "t-1"), &rec)
 	require.NoError(t, err)
 	require.True(t, found)
 	for _, c := range rec.Containers {
@@ -61,9 +61,9 @@ func TestReportTaskGPU_UnknownTaskNoop(t *testing.T) {
 // field, once ReportTaskGPU has merged them onto the task record.
 func TestDescribeTasks_GPUIDsPopulatedFromReport(t *testing.T) {
 	svc, _ := newTestService(t)
-	kv, err := svc.bucket(testAccountID)
+	kv, err := svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
-	require.NoError(t, putJSON(kv, TaskKey("web", "t-1"), &TaskRecord{
+	require.NoError(t, putJSON(t.Context(), kv, TaskKey("web", "t-1"), &TaskRecord{
 		TaskID: "t-1", ARN: TaskARN(svc.region, testAccountID, "web", "t-1"), Cluster: "web",
 		Containers: []ContainerState{{Name: "trainer", Status: "RUNNING", GPUIDs: []string{"GPU-aaa", "GPU-bbb"}}},
 	}))
@@ -83,9 +83,9 @@ func TestDescribeTasks_GPUIDsPopulatedFromReport(t *testing.T) {
 // RUNNING report (or ReportTaskGPU) previously set.
 func TestRecordTaskState_PreservesGPUIDsAcrossStoppedReport(t *testing.T) {
 	svc, _ := newTestService(t)
-	kv, err := svc.bucket(testAccountID)
+	kv, err := svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
-	require.NoError(t, putJSON(kv, TaskKey("web", "t-1"), &TaskRecord{
+	require.NoError(t, putJSON(t.Context(), kv, TaskKey("web", "t-1"), &TaskRecord{
 		TaskID: "t-1", ARN: TaskARN(svc.region, testAccountID, "web", "t-1"), Cluster: "web",
 		ContainerInstanceID: "i-1",
 		Containers:          []ContainerState{{Name: "trainer", Status: "PENDING"}},
@@ -99,7 +99,7 @@ func TestRecordTaskState_PreservesGPUIDsAcrossStoppedReport(t *testing.T) {
 	}))
 
 	var rec TaskRecord
-	found, err := getJSON(kv, TaskKey("web", "t-1"), &rec)
+	found, err := getJSON(t.Context(), kv, TaskKey("web", "t-1"), &rec)
 	require.NoError(t, err)
 	require.True(t, found)
 	require.Len(t, rec.Containers, 1)
@@ -114,7 +114,7 @@ func TestRecordTaskState_PreservesGPUIDsAcrossStoppedReport(t *testing.T) {
 		Containers: []bus.ContainerStatus{{Name: "trainer", Status: bus.TaskStatusStopped, ExitCode: &exit}},
 	}))
 
-	found, err = getJSON(kv, TaskKey("web", "t-1"), &rec)
+	found, err = getJSON(t.Context(), kv, TaskKey("web", "t-1"), &rec)
 	require.NoError(t, err)
 	require.True(t, found)
 	require.Len(t, rec.Containers, 1)
@@ -135,9 +135,9 @@ func TestRecordTaskState_PreservesGPUIDsAcrossStoppedReport(t *testing.T) {
 // pinned set normally (e.g. a corrected report, or GPUs reassigned on a retry).
 func TestRecordTaskState_ReportedGPUIDsOverwritePinned(t *testing.T) {
 	svc, _ := newTestService(t)
-	kv, err := svc.bucket(testAccountID)
+	kv, err := svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
-	require.NoError(t, putJSON(kv, TaskKey("web", "t-1"), &TaskRecord{
+	require.NoError(t, putJSON(t.Context(), kv, TaskKey("web", "t-1"), &TaskRecord{
 		TaskID: "t-1", ARN: TaskARN(svc.region, testAccountID, "web", "t-1"), Cluster: "web",
 		ContainerInstanceID: "i-1",
 		Containers:          []ContainerState{{Name: "trainer", Status: "RUNNING", GPUIDs: []string{"GPU-aaa"}}},
@@ -150,7 +150,7 @@ func TestRecordTaskState_ReportedGPUIDsOverwritePinned(t *testing.T) {
 	}))
 
 	var rec TaskRecord
-	found, err := getJSON(kv, TaskKey("web", "t-1"), &rec)
+	found, err := getJSON(t.Context(), kv, TaskKey("web", "t-1"), &rec)
 	require.NoError(t, err)
 	require.True(t, found)
 	require.Len(t, rec.Containers, 1)
@@ -161,9 +161,9 @@ func TestRecordTaskState_ReportedGPUIDsOverwritePinned(t *testing.T) {
 // nothing to preserve, and the fallback must not fabricate a value.
 func TestRecordTaskState_NonGPUContainerStaysEmpty(t *testing.T) {
 	svc, _ := newTestService(t)
-	kv, err := svc.bucket(testAccountID)
+	kv, err := svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
-	require.NoError(t, putJSON(kv, TaskKey("web", "t-1"), &TaskRecord{
+	require.NoError(t, putJSON(t.Context(), kv, TaskKey("web", "t-1"), &TaskRecord{
 		TaskID: "t-1", ARN: TaskARN(svc.region, testAccountID, "web", "t-1"), Cluster: "web",
 		ContainerInstanceID: "i-1",
 		Containers:          []ContainerState{{Name: "web", Status: "PENDING"}},
@@ -183,7 +183,7 @@ func TestRecordTaskState_NonGPUContainerStaysEmpty(t *testing.T) {
 	}))
 
 	var rec TaskRecord
-	found, err := getJSON(kv, TaskKey("web", "t-1"), &rec)
+	found, err := getJSON(t.Context(), kv, TaskKey("web", "t-1"), &rec)
 	require.NoError(t, err)
 	require.True(t, found)
 	require.Len(t, rec.Containers, 1)

spinifex/handlers/ecs/instance_recovery_test.go

@@ -24,10 +24,10 @@ func reRegister(t *testing.T, svc *Service, cluster, id string) {
 
 func instanceStatus(t *testing.T, svc *Service, cluster, id string) InstanceRecord {
 	t.Helper()
-	kv, err := svc.bucket(testAccountID)
+	kv, err := svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
 	var rec InstanceRecord
-	found, err := getJSON(kv, InstanceKey(cluster, id), &rec)
+	found, err := getJSON(t.Context(), kv, InstanceKey(cluster, id), &rec)
 	require.NoError(t, err)
 	require.True(t, found)
 	return rec
@@ -42,11 +42,11 @@ func TestRegister_RestoresReapedInstanceToActive(t *testing.T) {
 	require.NoError(t, err)
 	registerInstance(t, svc, "web", "i-1", 1024, 2048)
 
-	kv, err := svc.bucket(testAccountID)
+	kv, err := svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
 	rec := instanceStatus(t, svc, "web", "i-1")
 	rec.LastSeen = time.Now().UTC().Add(-2 * heartbeatTimeout)
-	require.NoError(t, putJSON(kv, InstanceKey("web", "i-1"), &rec))
+	require.NoError(t, putJSON(t.Context(), kv, InstanceKey("web", "i-1"), &rec))
 
 	NewScheduler(nc, svc, "test-holder").reapBucket(context.Background(), kv, testAccountID, time.Now().UTC())
 
@@ -83,10 +83,10 @@ func TestReaper_ReleasesCapacityAndRecoversClean(t *testing.T) {
 	require.Equal(t, 256, seeded.ReservedMemoryMiB)
 	require.Contains(t, seeded.PlacedTasks, taskID)
 
-	kv, err := svc.bucket(testAccountID)
+	kv, err := svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
 	seeded.LastSeen = time.Now().UTC().Add(-2 * heartbeatTimeout)
-	require.NoError(t, putJSON(kv, InstanceKey("web", "i-1"), &seeded))
+	require.NoError(t, putJSON(t.Context(), kv, InstanceKey("web", "i-1"), &seeded))
 
 	NewScheduler(nc, svc, "test-holder").reapBucket(context.Background(), kv, testAccountID, time.Now().UTC())
 

spinifex/handlers/ecs/instances.go

@@ -10,19 +10,19 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ecs"
 	"github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // listInstanceRecords returns all container-instance records in a cluster.
-func (s *Service) listInstanceRecords(kv nats.KeyValue, cluster string) ([]InstanceRecord, error) {
-	keys, err := keysWithPrefix(kv, InstancesPrefix(cluster))
+func (s *Service) listInstanceRecords(ctx context.Context, kv jetstream.KeyValue, cluster string) ([]InstanceRecord, error) {
+	keys, err := keysWithPrefix(ctx, kv, InstancesPrefix(cluster))
 	if err != nil {
 		return nil, err
 	}
 	out := make([]InstanceRecord, 0, len(keys))
 	for _, k := range keys {
 		var rec InstanceRecord
-		found, err := getJSON(kv, k, &rec)
+		found, err := getJSON(ctx, kv, k, &rec)
 		if err != nil {
 			return nil, err
 		}
@@ -36,17 +36,17 @@ func (s *Service) listInstanceRecords(kv nats.KeyValue, cluster string) ([]Insta
 // RegisterContainerInstance is the AWS-API registration path. In 4e the agent
 // normally registers over the Layer-2 bus; this keeps API parity by writing the
 // same record shape from an explicit call.
-func (s *Service) RegisterContainerInstance(_ context.Context, input *ecs.RegisterContainerInstanceInput, accountID string) (*ecs.RegisterContainerInstanceOutput, error) {
+func (s *Service) RegisterContainerInstance(ctx context.Context, input *ecs.RegisterContainerInstanceInput, accountID string) (*ecs.RegisterContainerInstanceOutput, error) {
 	cluster := clusterShortName(aws.StringValue(input.Cluster))
 	instanceID := aws.StringValue(input.InstanceIdentityDocument)
 	if instanceID == "" {
 		instanceID = "ci-" + time.Now().UTC().Format("20060102150405")
 	}
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
-	rec, err := s.upsertInstance(kv, accountID, cluster, instanceID, func(r *InstanceRecord) {
+	rec, err := s.upsertInstance(ctx, kv, accountID, cluster, instanceID, func(r *InstanceRecord) {
 		for _, res := range input.TotalResources {
 			switch aws.StringValue(res.Name) {
 			case "CPU":
@@ -79,9 +79,9 @@ func (s *Service) RegisterContainerInstance(_ context.Context, input *ecs.Regist
 }
 
 // DescribeContainerInstances returns records for the named container instances.
-func (s *Service) DescribeContainerInstances(_ context.Context, input *ecs.DescribeContainerInstancesInput, accountID string) (*ecs.DescribeContainerInstancesOutput, error) {
+func (s *Service) DescribeContainerInstances(ctx context.Context, input *ecs.DescribeContainerInstancesInput, accountID string) (*ecs.DescribeContainerInstancesOutput, error) {
 	cluster := clusterShortName(aws.StringValue(input.Cluster))
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -89,7 +89,7 @@ func (s *Service) DescribeContainerInstances(_ context.Context, input *ecs.Descr
 	for _, ref := range awsStringSlice(input.ContainerInstances) {
 		id := containerInstanceShortID(ref)
 		var rec InstanceRecord
-		found, err := getJSON(kv, InstanceKey(cluster, id), &rec)
+		found, err := getJSON(ctx, kv, InstanceKey(cluster, id), &rec)
 		if err != nil {
 			return nil, err
 		}
@@ -103,13 +103,13 @@ func (s *Service) DescribeContainerInstances(_ context.Context, input *ecs.Descr
 }
 
 // ListContainerInstances returns the ARNs of all container instances in a cluster.
-func (s *Service) ListContainerInstances(_ context.Context, input *ecs.ListContainerInstancesInput, accountID string) (*ecs.ListContainerInstancesOutput, error) {
+func (s *Service) ListContainerInstances(ctx context.Context, input *ecs.ListContainerInstancesInput, accountID string) (*ecs.ListContainerInstancesOutput, error) {
 	cluster := clusterShortName(aws.StringValue(input.Cluster))
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
-	recs, err := s.listInstanceRecords(kv, cluster)
+	recs, err := s.listInstanceRecords(ctx, kv, cluster)
 	if err != nil {
 		return nil, err
 	}
@@ -122,9 +122,9 @@ func (s *Service) ListContainerInstances(_ context.Context, input *ecs.ListConta
 
 // upsertInstance reads-or-creates the instance record, applies mutate, and writes
 // it back. Used by both the AWS register path and the bus register handler.
-func (s *Service) upsertInstance(kv nats.KeyValue, accountID, cluster, instanceID string, mutate func(*InstanceRecord)) (*InstanceRecord, error) {
+func (s *Service) upsertInstance(ctx context.Context, kv jetstream.KeyValue, accountID, cluster, instanceID string, mutate func(*InstanceRecord)) (*InstanceRecord, error) {
 	var rec InstanceRecord
-	found, err := getJSON(kv, InstanceKey(cluster, instanceID), &rec)
+	found, err := getJSON(ctx, kv, InstanceKey(cluster, instanceID), &rec)
 	if err != nil {
 		return nil, err
 	}
@@ -142,7 +142,7 @@ func (s *Service) upsertInstance(kv nats.KeyValue, accountID, cluster, instanceI
 	if mutate != nil {
 		mutate(&rec)
 	}
-	if err := putJSON(kv, InstanceKey(cluster, instanceID), &rec); err != nil {
+	if err := putJSON(ctx, kv, InstanceKey(cluster, instanceID), &rec); err != nil {
 		return nil, err
 	}
 	return &rec, nil
@@ -184,12 +184,12 @@ func (s *Service) instanceToAWS(r *InstanceRecord) *ecs.ContainerInstance {
 // --- Layer-2 bus event handlers (called by the scheduler) ---
 
 // recordRegister upserts a container-instance record from a bus RegisterInstance.
-func (s *Service) recordRegister(msg *bus.RegisterInstance) error {
-	kv, err := s.bucket(msg.AccountID)
+func (s *Service) recordRegister(ctx context.Context, msg *bus.RegisterInstance) error {
+	kv, err := s.bucket(ctx, msg.AccountID)
 	if err != nil {
 		return err
 	}
-	_, err = s.upsertInstance(kv, msg.AccountID, msg.ClusterName, msg.InstanceID, func(r *InstanceRecord) {
+	_, err = s.upsertInstance(ctx, kv, msg.AccountID, msg.ClusterName, msg.InstanceID, func(r *InstanceRecord) {
 		r.AZ = msg.AZ
 		r.Hostname = msg.Hostname
 		r.AgentVersion = msg.AgentVersion
@@ -203,13 +203,13 @@ func (s *Service) recordRegister(msg *bus.RegisterInstance) error {
 }
 
 // recordHeartbeat refreshes an instance's LastSeen and status from a bus beat.
-func (s *Service) recordHeartbeat(msg *bus.Heartbeat) error {
-	kv, err := s.bucket(msg.AccountID)
+func (s *Service) recordHeartbeat(ctx context.Context, msg *bus.Heartbeat) error {
+	kv, err := s.bucket(ctx, msg.AccountID)
 	if err != nil {
 		return err
 	}
 	var rec InstanceRecord
-	found, err := getJSON(kv, InstanceKey(msg.ClusterName, msg.InstanceID), &rec)
+	found, err := getJSON(ctx, kv, InstanceKey(msg.ClusterName, msg.InstanceID), &rec)
 	if err != nil || !found {
 		return err
 	}
@@ -217,18 +217,18 @@ func (s *Service) recordHeartbeat(msg *bus.Heartbeat) error {
 	if msg.Status != "" {
 		rec.Status = msg.Status
 	}
-	return putJSON(kv, InstanceKey(msg.ClusterName, msg.InstanceID), &rec)
+	return putJSON(ctx, kv, InstanceKey(msg.ClusterName, msg.InstanceID), &rec)
 }
 
 // recordTaskState applies an agent task-state report: it updates the task record
 // and, on STOPPED, releases the reserved capacity back to the instance.
 func (s *Service) recordTaskState(ctx context.Context, msg *bus.TaskState) error {
-	kv, err := s.bucket(msg.AccountID)
+	kv, err := s.bucket(ctx, msg.AccountID)
 	if err != nil {
 		return err
 	}
 	var task TaskRecord
-	found, err := getJSON(kv, TaskKey(msg.ClusterName, msg.TaskID), &task)
+	found, err := getJSON(ctx, kv, TaskKey(msg.ClusterName, msg.TaskID), &task)
 	if err != nil || !found {
 		return err
 	}
@@ -266,7 +266,7 @@ func (s *Service) recordTaskState(ctx context.Context, msg *bus.TaskState) error
 			task.StoppedReason = msg.Reason
 		}
 	}
-	if err := putJSON(kv, TaskKey(msg.ClusterName, msg.TaskID), &task); err != nil {
+	if err := putJSON(ctx, kv, TaskKey(msg.ClusterName, msg.TaskID), &task); err != nil {
 		return err
 	}
 
@@ -287,13 +287,13 @@ func (s *Service) recordTaskState(ctx context.Context, msg *bus.TaskState) error
 		}
 		s.deregisterServiceTargets(ctx, kv, msg.AccountID, &task)
 		s.releaseTaskPublicIP(ctx, msg.AccountID, &task)
-		s.reclaimAssignInbox(kv, msg.ClusterName, task.ContainerInstanceID, msg.TaskID)
-		s.reclaimStopInbox(kv, msg.ClusterName, task.ContainerInstanceID, msg.TaskID)
+		s.reclaimAssignInbox(ctx, kv, msg.ClusterName, task.ContainerInstanceID, msg.TaskID)
+		s.reclaimStopInbox(ctx, kv, msg.ClusterName, task.ContainerInstanceID, msg.TaskID)
 		s.reclaimTaskENI(ctx, msg.AccountID, &task)
-		if perr := putJSON(kv, TaskKey(msg.ClusterName, msg.TaskID), &task); perr != nil {
+		if perr := putJSON(ctx, kv, TaskKey(msg.ClusterName, msg.TaskID), &task); perr != nil {
 			slog.ErrorContext(ctx, "ECS task STOPPED: persist after EIP release failed", "task", msg.TaskID, "err", perr)
 		}
-		return s.releaseReservation(kv, msg.ClusterName, task.ContainerInstanceID, msg.TaskID, task.ReservedCPU, task.ReservedMemoryMiB, task.GPU)
+		return s.releaseReservation(ctx, kv, msg.ClusterName, task.ContainerInstanceID, msg.TaskID, task.ReservedCPU, task.ReservedMemoryMiB, task.GPU)
 	}
 	return nil
 }
@@ -332,14 +332,14 @@ func (s *Service) SubmitTaskStateChange(ctx context.Context, input *ecs.SubmitTa
 // recordDeploymentFailure increments the failed-task counter on the deployment
 // that launched a task which stopped before ever running, driving the service's
 // deployment circuit breaker. No-op for a non-service task or an unknown deployment.
-func (s *Service) recordDeploymentFailure(ctx context.Context, kv nats.KeyValue, cluster string, task *TaskRecord) {
+func (s *Service) recordDeploymentFailure(ctx context.Context, kv jetstream.KeyValue, cluster string, task *TaskRecord) {
 	name := serviceNameFromGroup(task.Group)
 	depID := deploymentIDFromStartedBy(task.StartedBy)
 	if name == "" || depID == "" {
 		return
 	}
 	var svc ServiceRecord
-	found, err := getJSON(kv, ServiceKey(cluster, name), &svc)
+	found, err := getJSON(ctx, kv, ServiceKey(cluster, name), &svc)
 	if err != nil || !found {
 		return
 	}
@@ -347,7 +347,7 @@ func (s *Service) recordDeploymentFailure(ctx context.Context, kv nats.KeyValue,
 		if svc.Deployments[i].ID == depID {
 			svc.Deployments[i].FailedTasks++
 			svc.Deployments[i].UpdatedAt = time.Now().UTC()
-			if perr := putJSON(kv, ServiceKey(cluster, name), &svc); perr != nil {
+			if perr := putJSON(ctx, kv, ServiceKey(cluster, name), &svc); perr != nil {
 				slog.ErrorContext(ctx, "ECS deployment failure accounting: persist failed", "service", name, "err", perr)
 			}
 			return
@@ -356,9 +356,9 @@ func (s *Service) recordDeploymentFailure(ctx context.Context, kv nats.KeyValue,
 }
 
 // releaseReservation returns a stopped task's capacity to its instance under CAS.
-func (s *Service) releaseReservation(kv nats.KeyValue, cluster, instanceID, taskID string, cpu, mem, gpu int) error {
+func (s *Service) releaseReservation(ctx context.Context, kv jetstream.KeyValue, cluster, instanceID, taskID string, cpu, mem, gpu int) error {
 	for range reservePlacementRetries {
-		entry, err := kv.Get(InstanceKey(cluster, instanceID))
+		entry, err := kv.Get(ctx, InstanceKey(cluster, instanceID))
 		if err != nil {
 			return nil //nolint:nilerr // instance gone; nothing to release
 		}
@@ -374,7 +374,7 @@ func (s *Service) releaseReservation(kv nats.KeyValue, cluster, instanceID, task
 		if merr != nil {
 			return merr
 		}
-		if _, uerr := kv.Update(InstanceKey(cluster, instanceID), data, entry.Revision()); uerr == nil {
+		if _, uerr := kv.Update(ctx, InstanceKey(cluster, instanceID), data, entry.Revision()); uerr == nil {
 			return nil
 		}
 	}

spinifex/handlers/ecs/poll.go

@@ -4,7 +4,7 @@ import (
 	"context"
 
 	"github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // PollAssignmentsInput is the agent's poll request. ContainerInstance names the
@@ -34,10 +34,10 @@ type PollAssignmentsOutput struct {
 // writer of the inbox; the agent is the only reader/acker. Account is
 // authoritative from accountID (SigV4 caller), so an instance cannot drain
 // another account's inbox.
-func (s *Service) PollAssignments(_ context.Context, input *PollAssignmentsInput, accountID string) (*PollAssignmentsOutput, error) {
+func (s *Service) PollAssignments(ctx context.Context, input *PollAssignmentsInput, accountID string) (*PollAssignmentsOutput, error) {
 	cluster := clusterShortName(input.Cluster)
 	instanceID := containerInstanceShortID(input.ContainerInstance)
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -46,23 +46,23 @@ func (s *Service) PollAssignments(_ context.Context, input *PollAssignmentsInput
 		if taskID == "" {
 			continue
 		}
-		_ = kv.Delete(AssignmentKey(cluster, instanceID, taskShortID(taskID)))
+		_ = kv.Delete(ctx, AssignmentKey(cluster, instanceID, taskShortID(taskID)))
 	}
 	for _, taskID := range input.AckStopIDs {
 		if taskID == "" {
 			continue
 		}
-		_ = kv.Delete(StopKey(cluster, instanceID, taskShortID(taskID)))
+		_ = kv.Delete(ctx, StopKey(cluster, instanceID, taskShortID(taskID)))
 	}
 
-	keys, err := keysWithPrefix(kv, AssignmentsPrefix(cluster, instanceID))
+	keys, err := keysWithPrefix(ctx, kv, AssignmentsPrefix(cluster, instanceID))
 	if err != nil {
 		return nil, err
 	}
 	out := &PollAssignmentsOutput{}
 	for _, key := range keys {
 		var as bus.Assign
-		found, err := getJSON(kv, key, &as)
+		found, err := getJSON(ctx, kv, key, &as)
 		if err != nil {
 			return nil, err
 		}
@@ -71,13 +71,13 @@ func (s *Service) PollAssignments(_ context.Context, input *PollAssignmentsInput
 		}
 	}
 
-	stopKeys, err := keysWithPrefix(kv, StopsPrefix(cluster, instanceID))
+	stopKeys, err := keysWithPrefix(ctx, kv, StopsPrefix(cluster, instanceID))
 	if err != nil {
 		return nil, err
 	}
 	for _, key := range stopKeys {
 		var sd bus.StopDirective
-		found, err := getJSON(kv, key, &sd)
+		found, err := getJSON(ctx, kv, key, &sd)
 		if err != nil {
 			return nil, err
 		}
@@ -90,18 +90,18 @@ func (s *Service) PollAssignments(_ context.Context, input *PollAssignmentsInput
 
 // reclaimAssignInbox removes a task's inbox entry on the STOPPED transition so a
 // stopped task is never re-delivered to the agent. Best-effort.
-func (s *Service) reclaimAssignInbox(kv nats.KeyValue, cluster, instanceID, taskID string) {
+func (s *Service) reclaimAssignInbox(ctx context.Context, kv jetstream.KeyValue, cluster, instanceID, taskID string) {
 	if instanceID == "" || taskID == "" {
 		return
 	}
-	_ = kv.Delete(AssignmentKey(cluster, instanceID, taskID))
+	_ = kv.Delete(ctx, AssignmentKey(cluster, instanceID, taskID))
 }
 
 // reclaimStopInbox removes a task's stop-inbox entry on the STOPPED transition so
 // a reaped task is never re-delivered as a stop. Best-effort.
-func (s *Service) reclaimStopInbox(kv nats.KeyValue, cluster, instanceID, taskID string) {
+func (s *Service) reclaimStopInbox(ctx context.Context, kv jetstream.KeyValue, cluster, instanceID, taskID string) {
 	if instanceID == "" || taskID == "" {
 		return
 	}
-	_ = kv.Delete(StopKey(cluster, instanceID, taskID))
+	_ = kv.Delete(ctx, StopKey(cluster, instanceID, taskID))
 }

spinifex/handlers/ecs/poll_test.go

@@ -16,9 +16,9 @@ import (
 func TestPollAssignments_DrainsStopsAndAcks(t *testing.T) {
 	svc, _, kv := serviceTestRig(t)
 
-	require.NoError(t, putJSON(kv, AssignmentKey("web", "i-1", "t-1"),
+	require.NoError(t, putJSON(t.Context(), kv, AssignmentKey("web", "i-1", "t-1"),
 		&bus.Assign{TaskID: "t-1", ClusterName: "web", InstanceID: "i-1"}))
-	require.NoError(t, putJSON(kv, StopKey("web", "i-1", "t-2"),
+	require.NoError(t, putJSON(t.Context(), kv, StopKey("web", "i-1", "t-2"),
 		&bus.StopDirective{TaskID: "t-2", Reason: "drain"}))
 
 	out, err := svc.PollAssignments(context.Background(), &PollAssignmentsInput{
@@ -41,7 +41,7 @@ func TestPollAssignments_DrainsStopsAndAcks(t *testing.T) {
 	assert.Empty(t, out.Assignments)
 	assert.Empty(t, out.Stops)
 
-	found, err := getJSON(kv, StopKey("web", "i-1", "t-2"), &bus.StopDirective{})
+	found, err := getJSON(t.Context(), kv, StopKey("web", "i-1", "t-2"), &bus.StopDirective{})
 	require.NoError(t, err)
 	assert.False(t, found)
 }

spinifex/handlers/ecs/scheduler.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Scheduler timing. The leader lease lives in spinifex-ecs-leader (60s TTL); the
@@ -70,21 +71,27 @@ func (sc *Scheduler) Run(ctx context.Context) {
 		case <-leaseTicker.C:
 			sc.evaluateLeadership(ctx)
 		case <-reaperTicker.C:
-			if sc.isLeader() {
-				sc.reap()
-			}
+			sc.runIfLeader("instance reap", func() error { return sc.reap(ctx) })
 		case <-reconcileTicker.C:
-			if sc.isLeader() {
-				sc.svc.reconcileAllServices()
-			}
+			sc.runIfLeader("service reconcile", func() error { return sc.svc.reconcileAllServices(ctx) })
 		case <-sweepTicker.C:
-			if sc.isLeader() {
-				sc.sweepStoppedTasks()
-			}
+			sc.runIfLeader("stopped-task sweep", func() error { return sc.sweepStoppedTasks(ctx) })
 		}
 	}
 }
 
+// runIfLeader runs one leader-only pass, logging a pass that could not complete.
+// A pass errors only when it could not observe the whole fleet, so this log is
+// the operator's signal that a tick was skipped rather than found nothing to do.
+func (sc *Scheduler) runIfLeader(pass string, run func() error) {
+	if !sc.isLeader() {
+		return
+	}
+	if err := run(); err != nil {
+		slog.Error("ECS scheduler: pass failed", "pass", pass, "err", err)
+	}
+}
+
 func (sc *Scheduler) isLeader() bool {
 	sc.mu.Lock()
 	defer sc.mu.Unlock()
@@ -94,7 +101,7 @@ func (sc *Scheduler) isLeader() bool {
 // evaluateLeadership acquires or refreshes the lease, wiring up (or tearing down)
 // the bus subscriptions as leadership changes.
 func (sc *Scheduler) evaluateLeadership(ctx context.Context) {
-	won := sc.acquireOrRefresh()
+	won := sc.acquireOrRefresh(ctx)
 	sc.mu.Lock()
 	was := sc.leader
 	sc.leader = won
@@ -115,19 +122,19 @@ func (sc *Scheduler) evaluateLeadership(ctx context.Context) {
 
 // acquireOrRefresh tries to claim the scheduler lease, refreshing it (resetting
 // the TTL) when we already hold it. Returns true when we are the leader.
-func (sc *Scheduler) acquireOrRefresh() bool {
-	js, err := sc.nc.JetStream()
+func (sc *Scheduler) acquireOrRefresh(ctx context.Context) bool {
+	js, err := jetstream.New(sc.nc)
 	if err != nil {
 		return false
 	}
-	kv, err := InitLeaderBucket(js)
+	kv, err := InitLeaderBucket(ctx, js)
 	if err != nil {
 		return false
 	}
-	if _, err := kv.Create(schedulerLeaderKey, []byte(sc.holder)); err == nil {
+	if _, err := kv.Create(ctx, schedulerLeaderKey, []byte(sc.holder)); err == nil {
 		return true
 	}
-	entry, err := kv.Get(schedulerLeaderKey)
+	entry, err := kv.Get(ctx, schedulerLeaderKey)
 	if err != nil {
 		return false
 	}
@@ -135,7 +142,7 @@ func (sc *Scheduler) acquireOrRefresh() bool {
 		return false // someone else holds it
 	}
 	// We hold it — refresh to reset the TTL.
-	if _, err := kv.Put(schedulerLeaderKey, []byte(sc.holder)); err != nil {
+	if _, err := kv.Put(ctx, schedulerLeaderKey, []byte(sc.holder)); err != nil {
 		return false
 	}
 	return true
@@ -144,16 +151,20 @@ func (sc *Scheduler) acquireOrRefresh() bool {
 // relinquish drops subscriptions and deletes the lease key on shutdown.
 func (sc *Scheduler) relinquish() {
 	sc.unsubscribeBus()
-	js, err := sc.nc.JetStream()
+	// Run's ctx is already cancelled by the time it calls relinquish, so the
+	// release runs on its own context — a captured ctx would fail the delete and
+	// leave the lease for the TTL to reap.
+	ctx := context.Background()
+	js, err := jetstream.New(sc.nc)
 	if err != nil {
 		return
 	}
-	kv, err := InitLeaderBucket(js)
+	kv, err := InitLeaderBucket(ctx, js)
 	if err != nil {
 		return
 	}
-	if entry, gerr := kv.Get(schedulerLeaderKey); gerr == nil && string(entry.Value()) == sc.holder {
-		_ = kv.Delete(schedulerLeaderKey)
+	if entry, gerr := kv.Get(ctx, schedulerLeaderKey); gerr == nil && string(entry.Value()) == sc.holder {
+		_ = kv.Delete(ctx, schedulerLeaderKey)
 	}
 }
 
@@ -197,7 +208,9 @@ func (sc *Scheduler) onRegister(msg *nats.Msg) {
 		slog.Warn("ECS scheduler: bad register payload", "err", err)
 		return
 	}
-	if err := sc.svc.recordRegister(&m); err != nil {
+	// A bus callback outlives the Run context that wired its subscription, so it
+	// binds its own — a cancelled ctx would abandon the KV write mid-record.
+	if err := sc.svc.recordRegister(context.Background(), &m); err != nil {
 		slog.Error("ECS scheduler: record register failed", "instance", m.InstanceID, "err", err)
 		return
 	}
@@ -209,7 +222,7 @@ func (sc *Scheduler) onHeartbeat(msg *nats.Msg) {
 	if err := json.Unmarshal(msg.Data, &m); err != nil {
 		return
 	}
-	if err := sc.svc.recordHeartbeat(&m); err != nil {
+	if err := sc.svc.recordHeartbeat(context.Background(), &m); err != nil {
 		slog.Debug("ECS scheduler: record heartbeat failed", "instance", m.InstanceID, "err", err)
 	}
 }
@@ -227,29 +240,33 @@ func (sc *Scheduler) onTaskState(msg *nats.Msg) {
 }
 
 // reap marks instances that have missed their heartbeat window DRAINING and stops
-// their tasks, releasing capacity. Iterates every ECS account bucket.
-func (sc *Scheduler) reap() {
-	ctx := context.Background()
-	js, err := sc.nc.JetStream()
+// their tasks, releasing capacity. Iterates every ECS account bucket. Returns an
+// error when the account enumeration could not be completed, so a pass that
+// could not see the whole fleet is reported rather than read as "nothing to
+// reap" — every unlisted account keeps its dead instances' capacity held.
+func (sc *Scheduler) reap(ctx context.Context) error {
+	js, err := jetstream.New(sc.nc)
 	if err != nil {
-		return
+		return err
+	}
+	buckets, err := accountBuckets(ctx, sc.nc)
+	if err != nil {
+		return err
 	}
 	now := time.Now().UTC()
-	for bucket := range js.KeyValueStoreNames() {
-		accountID, ok := accountIDFromBucket(bucket)
-		if !ok {
-			continue
-		}
-		kv, err := js.KeyValue(bucket)
+	for _, bucket := range buckets {
+		kv, err := js.KeyValue(ctx, bucket.name)
 		if err != nil {
+			slog.Error("ECS scheduler: open bucket failed", "bucket", bucket.name, "err", err)
 			continue
 		}
-		sc.reapBucket(ctx, kv, accountID, now)
+		sc.reapBucket(ctx, kv, bucket.accountID, now)
 	}
+	return nil
 }
 
-func (sc *Scheduler) reapBucket(ctx context.Context, kv nats.KeyValue, accountID string, now time.Time) {
-	keys, err := keysWithPrefix(kv, "clusters/")
+func (sc *Scheduler) reapBucket(ctx context.Context, kv jetstream.KeyValue, accountID string, now time.Time) {
+	keys, err := keysWithPrefix(ctx, kv, "clusters/")
 	if err != nil {
 		return
 	}
@@ -258,7 +275,7 @@ func (sc *Scheduler) reapBucket(ctx context.Context, kv nats.KeyValue, accountID
 			continue
 		}
 		var inst InstanceRecord
-		found, err := getJSON(kv, k, &inst)
+		found, err := getJSON(ctx, kv, k, &inst)
 		if err != nil || !found {
 			continue
 		}
@@ -269,7 +286,7 @@ func (sc *Scheduler) reapBucket(ctx context.Context, kv nats.KeyValue, accountID
 			"lastSeen", inst.LastSeen.Format(time.RFC3339))
 		inst.Status = InstanceStatusDraining
 		inst.Reaped = true
-		if perr := putJSON(kv, k, &inst); perr != nil {
+		if perr := putJSON(ctx, kv, k, &inst); perr != nil {
 			continue
 		}
 		sc.stopInstanceTasks(ctx, kv, accountID, inst.Cluster, inst.InstanceID)
@@ -278,14 +295,14 @@ func (sc *Scheduler) reapBucket(ctx context.Context, kv nats.KeyValue, accountID
 
 // stopInstanceTasks transitions a reaped instance's non-stopped tasks to STOPPED
 // and reclaims each awsvpc task's ENI (leak guard for a dead agent).
-func (sc *Scheduler) stopInstanceTasks(ctx context.Context, kv nats.KeyValue, accountID, cluster, instanceID string) {
-	keys, err := keysWithPrefix(kv, TasksPrefix(cluster))
+func (sc *Scheduler) stopInstanceTasks(ctx context.Context, kv jetstream.KeyValue, accountID, cluster, instanceID string) {
+	keys, err := keysWithPrefix(ctx, kv, TasksPrefix(cluster))
 	if err != nil {
 		return
 	}
 	for _, k := range keys {
 		var task TaskRecord
-		found, err := getJSON(kv, k, &task)
+		found, err := getJSON(ctx, kv, k, &task)
 		if err != nil || !found {
 			continue
 		}

spinifex/handlers/ecs/service.go

@@ -18,6 +18,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/iam"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // ECSService is the ECS control-plane contract implemented by the daemon and
@@ -151,26 +152,26 @@ func NewService(nc *nats.Conn, region, suffix string) *Service {
 // defaultCluster is the implicit cluster name when a request omits one.
 const defaultCluster = "default"
 
-func (s *Service) js() (nats.JetStreamContext, error) {
+func (s *Service) js() (jetstream.JetStream, error) {
 	if s.nc == nil {
 		return nil, errors.New("ecs service: nil nats connection")
 	}
-	return s.nc.JetStream()
+	return jetstream.New(s.nc)
 }
 
 // bucket returns the per-account KV handle, creating it on first use.
-func (s *Service) bucket(accountID string) (nats.KeyValue, error) {
+func (s *Service) bucket(ctx context.Context, accountID string) (jetstream.KeyValue, error) {
 	js, err := s.js()
 	if err != nil {
 		return nil, err
 	}
-	return GetOrCreateAccountBucket(js, accountID)
+	return GetOrCreateAccountBucket(ctx, js, accountID)
 }
 
 // getJSON reads key into out. Returns (false, nil) when the key is absent.
-func getJSON(kv nats.KeyValue, key string, out any) (bool, error) {
-	entry, err := kv.Get(key)
-	if errors.Is(err, nats.ErrKeyNotFound) {
+func getJSON(ctx context.Context, kv jetstream.KeyValue, key string, out any) (bool, error) {
+	entry, err := kv.Get(ctx, key)
+	if errors.Is(err, jetstream.ErrKeyNotFound) {
 		return false, nil
 	}
 	if err != nil {
@@ -183,12 +184,12 @@ func getJSON(kv nats.KeyValue, key string, out any) (bool, error) {
 }
 
 // putJSON marshals v and writes it at key.
-func putJSON(kv nats.KeyValue, key string, v any) error {
+func putJSON(ctx context.Context, kv jetstream.KeyValue, key string, v any) error {
 	data, err := json.Marshal(v)
 	if err != nil {
 		return err
 	}
-	if _, err := kv.Put(key, data); err != nil {
+	if _, err := kv.Put(ctx, key, data); err != nil {
 		return fmt.Errorf("put %s: %w", key, err)
 	}
 	return nil
@@ -198,14 +199,14 @@ func putJSON(kv nats.KeyValue, key string, v any) error {
 // to sweep a cluster's whole key subtree (meta/instances/tasks/services/
 // assignments) after its tasks are stopped. Best-effort per key; the first
 // delete error is returned after attempting the rest.
-func deleteKeysWithPrefix(kv nats.KeyValue, prefix string) error {
-	keys, err := keysWithPrefix(kv, prefix)
+func deleteKeysWithPrefix(ctx context.Context, kv jetstream.KeyValue, prefix string) error {
+	keys, err := keysWithPrefix(ctx, kv, prefix)
 	if err != nil {
 		return err
 	}
 	var firstErr error
 	for _, k := range keys {
-		if derr := kv.Delete(k); derr != nil && firstErr == nil {
+		if derr := kv.Delete(ctx, k); derr != nil && firstErr == nil {
 			firstErr = derr
 		}
 	}
@@ -213,9 +214,9 @@ func deleteKeysWithPrefix(kv nats.KeyValue, prefix string) error {
 }
 
 // keysWithPrefix returns all KV keys under prefix.
-func keysWithPrefix(kv nats.KeyValue, prefix string) ([]string, error) {
-	keys, err := kv.Keys()
-	if errors.Is(err, nats.ErrNoKeysFound) {
+func keysWithPrefix(ctx context.Context, kv jetstream.KeyValue, prefix string) ([]string, error) {
+	keys, err := kv.Keys(ctx)
+	if errors.Is(err, jetstream.ErrNoKeysFound) {
 		return nil, nil
 	}
 	if err != nil {
@@ -235,18 +236,18 @@ func keysWithPrefix(kv nats.KeyValue, prefix string) ([]string, error) {
 
 // CreateCluster persists a cluster meta record. Idempotent: re-creating an
 // existing cluster returns the existing ACTIVE record.
-func (s *Service) CreateCluster(_ context.Context, input *ecs.CreateClusterInput, accountID string) (*ecs.CreateClusterOutput, error) {
+func (s *Service) CreateCluster(ctx context.Context, input *ecs.CreateClusterInput, accountID string) (*ecs.CreateClusterOutput, error) {
 	name := aws.StringValue(input.ClusterName)
 	if name == "" {
 		name = defaultCluster
 	}
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
 	var rec ClusterRecord
-	found, err := getJSON(kv, ClusterMetaKey(name), &rec)
+	found, err := getJSON(ctx, kv, ClusterMetaKey(name), &rec)
 	if err != nil {
 		return nil, err
 	}
@@ -263,7 +264,7 @@ func (s *Service) CreateCluster(_ context.Context, input *ecs.CreateClusterInput
 				rec.Tags[aws.StringValue(t.Key)] = aws.StringValue(t.Value)
 			}
 		}
-		if err := putJSON(kv, ClusterMetaKey(name), &rec); err != nil {
+		if err := putJSON(ctx, kv, ClusterMetaKey(name), &rec); err != nil {
 			return nil, err
 		}
 	}
@@ -272,8 +273,8 @@ func (s *Service) CreateCluster(_ context.Context, input *ecs.CreateClusterInput
 
 // DescribeClusters returns meta for the named clusters; unknown names are
 // silently skipped (AWS returns them under "failures", omitted here in v1).
-func (s *Service) DescribeClusters(_ context.Context, input *ecs.DescribeClustersInput, accountID string) (*ecs.DescribeClustersOutput, error) {
-	kv, err := s.bucket(accountID)
+func (s *Service) DescribeClusters(ctx context.Context, input *ecs.DescribeClustersInput, accountID string) (*ecs.DescribeClustersOutput, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -285,7 +286,7 @@ func (s *Service) DescribeClusters(_ context.Context, input *ecs.DescribeCluster
 	for _, name := range names {
 		name = clusterShortName(name)
 		var rec ClusterRecord
-		found, err := getJSON(kv, ClusterMetaKey(name), &rec)
+		found, err := getJSON(ctx, kv, ClusterMetaKey(name), &rec)
 		if err != nil {
 			return nil, err
 		}
@@ -297,12 +298,12 @@ func (s *Service) DescribeClusters(_ context.Context, input *ecs.DescribeCluster
 }
 
 // ListClusters returns the ARNs of all clusters in the account.
-func (s *Service) ListClusters(_ context.Context, _ *ecs.ListClustersInput, accountID string) (*ecs.ListClustersOutput, error) {
-	kv, err := s.bucket(accountID)
+func (s *Service) ListClusters(ctx context.Context, _ *ecs.ListClustersInput, accountID string) (*ecs.ListClustersOutput, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
-	keys, err := keysWithPrefix(kv, "clusters/")
+	keys, err := keysWithPrefix(ctx, kv, "clusters/")
 	if err != nil {
 		return nil, err
 	}
@@ -312,7 +313,7 @@ func (s *Service) ListClusters(_ context.Context, _ *ecs.ListClustersInput, acco
 			continue
 		}
 		var rec ClusterRecord
-		found, err := getJSON(kv, k, &rec)
+		found, err := getJSON(ctx, kv, k, &rec)
 		if err != nil {
 			return nil, err
 		}
@@ -384,12 +385,12 @@ func (s *Service) RegisterTaskDefinition(ctx context.Context, input *ecs.Registe
 		return nil, err
 	}
 	warnUnsupportedLogDrivers(ctx, family, input.ContainerDefinitions)
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
-	rev, err := s.nextRevision(kv, family)
+	rev, err := s.nextRevision(ctx, kv, family)
 	if err != nil {
 		return nil, err
 	}
@@ -408,10 +409,10 @@ func (s *Service) RegisterTaskDefinition(ctx context.Context, input *ecs.Registe
 		RegisteredAt:     time.Now().UTC(),
 		Containers:       containerDefsFromAWS(input.ContainerDefinitions),
 	}
-	if err := putJSON(kv, TaskDefRevKey(family, rev), &rec); err != nil {
+	if err := putJSON(ctx, kv, TaskDefRevKey(family, rev), &rec); err != nil {
 		return nil, err
 	}
-	if err := putJSON(kv, TaskDefLatestRevKey(family), rev); err != nil {
+	if err := putJSON(ctx, kv, TaskDefLatestRevKey(family), rev); err != nil {
 		return nil, err
 	}
 	return &ecs.RegisterTaskDefinitionOutput{TaskDefinition: rec.toAWS(), Tags: tagsToAWS(rec.Tags)}, nil
@@ -447,9 +448,9 @@ func warnUnsupportedLogDrivers(ctx context.Context, family string, defs []*ecs.C
 }
 
 // nextRevision reads the family's latest-rev and returns latest+1 (1 if absent).
-func (s *Service) nextRevision(kv nats.KeyValue, family string) (int, error) {
+func (s *Service) nextRevision(ctx context.Context, kv jetstream.KeyValue, family string) (int, error) {
 	var latest int
-	found, err := getJSON(kv, TaskDefLatestRevKey(family), &latest)
+	found, err := getJSON(ctx, kv, TaskDefLatestRevKey(family), &latest)
 	if err != nil {
 		return 0, err
 	}
@@ -462,17 +463,17 @@ func (s *Service) nextRevision(kv nats.KeyValue, family string) (int, error) {
 // DeregisterTaskDefinition marks a specific task-definition revision INACTIVE.
 // AWS requires an explicit family:revision (a bare family is rejected); the
 // revision stays describable, matching AWS. Idempotent.
-func (s *Service) DeregisterTaskDefinition(_ context.Context, input *ecs.DeregisterTaskDefinitionInput, accountID string) (*ecs.DeregisterTaskDefinitionOutput, error) {
+func (s *Service) DeregisterTaskDefinition(ctx context.Context, input *ecs.DeregisterTaskDefinitionInput, accountID string) (*ecs.DeregisterTaskDefinitionOutput, error) {
 	family, rev := parseTaskDefRef(aws.StringValue(input.TaskDefinition))
 	if family == "" || rev == 0 {
 		return nil, errors.New(awserrors.ErrorECSInvalidParameter)
 	}
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 	var rec TaskDefRecord
-	found, err := getJSON(kv, TaskDefRevKey(family, rev), &rec)
+	found, err := getJSON(ctx, kv, TaskDefRevKey(family, rev), &rec)
 	if err != nil {
 		return nil, err
 	}
@@ -480,19 +481,19 @@ func (s *Service) DeregisterTaskDefinition(_ context.Context, input *ecs.Deregis
 		return nil, errors.New(awserrors.ErrorECSInvalidParameter)
 	}
 	rec.Status = TaskDefStatusInactive
-	if err := putJSON(kv, TaskDefRevKey(family, rev), &rec); err != nil {
+	if err := putJSON(ctx, kv, TaskDefRevKey(family, rev), &rec); err != nil {
 		return nil, err
 	}
 	return &ecs.DeregisterTaskDefinitionOutput{TaskDefinition: rec.toAWS()}, nil
 }
 
 // DescribeTaskDefinition resolves "family", "family:rev" or an ARN to a revision.
-func (s *Service) DescribeTaskDefinition(_ context.Context, input *ecs.DescribeTaskDefinitionInput, accountID string) (*ecs.DescribeTaskDefinitionOutput, error) {
-	kv, err := s.bucket(accountID)
+func (s *Service) DescribeTaskDefinition(ctx context.Context, input *ecs.DescribeTaskDefinitionInput, accountID string) (*ecs.DescribeTaskDefinitionOutput, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
-	rec, err := s.resolveTaskDef(kv, aws.StringValue(input.TaskDefinition))
+	rec, err := s.resolveTaskDef(ctx, kv, aws.StringValue(input.TaskDefinition))
 	if err != nil {
 		return nil, err
 	}
@@ -502,8 +503,8 @@ func (s *Service) DescribeTaskDefinition(_ context.Context, input *ecs.DescribeT
 // ListTaskDefinitions returns revision ARNs, optionally filtered by family and
 // by status. Matching AWS, the status defaults to ACTIVE when unset, so
 // deregistered (INACTIVE) revisions drop off the default listing.
-func (s *Service) ListTaskDefinitions(_ context.Context, input *ecs.ListTaskDefinitionsInput, accountID string) (*ecs.ListTaskDefinitionsOutput, error) {
-	kv, err := s.bucket(accountID)
+func (s *Service) ListTaskDefinitions(ctx context.Context, input *ecs.ListTaskDefinitionsInput, accountID string) (*ecs.ListTaskDefinitionsOutput, error) {
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -515,7 +516,7 @@ func (s *Service) ListTaskDefinitions(_ context.Context, input *ecs.ListTaskDefi
 	if fam := aws.StringValue(input.FamilyPrefix); fam != "" {
 		prefix = TaskDefRevsPrefix(fam)
 	}
-	keys, err := keysWithPrefix(kv, prefix)
+	keys, err := keysWithPrefix(ctx, kv, prefix)
 	if err != nil {
 		return nil, err
 	}
@@ -525,7 +526,7 @@ func (s *Service) ListTaskDefinitions(_ context.Context, input *ecs.ListTaskDefi
 			continue
 		}
 		var rec TaskDefRecord
-		found, err := getJSON(kv, k, &rec)
+		found, err := getJSON(ctx, kv, k, &rec)
 		if err != nil {
 			return nil, err
 		}
@@ -538,14 +539,14 @@ func (s *Service) ListTaskDefinitions(_ context.Context, input *ecs.ListTaskDefi
 
 // resolveTaskDef loads the TaskDefRecord named by ref ("family", "family:rev",
 // or a task-definition ARN). A bare family resolves to its latest revision.
-func (s *Service) resolveTaskDef(kv nats.KeyValue, ref string) (*TaskDefRecord, error) {
+func (s *Service) resolveTaskDef(ctx context.Context, kv jetstream.KeyValue, ref string) (*TaskDefRecord, error) {
 	family, rev := parseTaskDefRef(ref)
 	if family == "" {
 		return nil, errors.New(awserrors.ErrorECSInvalidParameter)
 	}
 	if rev == 0 {
 		var latest int
-		found, err := getJSON(kv, TaskDefLatestRevKey(family), &latest)
+		found, err := getJSON(ctx, kv, TaskDefLatestRevKey(family), &latest)
 		if err != nil {
 			return nil, err
 		}
@@ -555,7 +556,7 @@ func (s *Service) resolveTaskDef(kv nats.KeyValue, ref string) (*TaskDefRecord,
 		rev = latest
 	}
 	var rec TaskDefRecord
-	found, err := getJSON(kv, TaskDefRevKey(family, rev), &rec)
+	found, err := getJSON(ctx, kv, TaskDefRevKey(family, rev), &rec)
 	if err != nil {
 		return nil, err
 	}

spinifex/handlers/ecs/service_test.go

@@ -166,7 +166,7 @@ func TestService_DescribeTaskDefinition_Unknown(t *testing.T) {
 // register path (the way the agent registers in production).
 func registerInstance(t *testing.T, svc *Service, cluster, id string, cpu, mem int) {
 	t.Helper()
-	require.NoError(t, svc.recordRegister(&bus.RegisterInstance{
+	require.NoError(t, svc.recordRegister(t.Context(), &bus.RegisterInstance{
 		AccountID:   testAccountID,
 		ClusterName: cluster,
 		InstanceID:  id,
@@ -393,7 +393,7 @@ func TestService_SubmitTaskStateChange_StopsTask(t *testing.T) {
 func TestService_RecordHeartbeat_UpdatesStatus(t *testing.T) {
 	svc, _ := newTestService(t)
 	registerInstance(t, svc, "web", "i-1", 1024, 2048)
-	require.NoError(t, svc.recordHeartbeat(&bus.Heartbeat{
+	require.NoError(t, svc.recordHeartbeat(t.Context(), &bus.Heartbeat{
 		AccountID: testAccountID, ClusterName: "web", InstanceID: "i-1", Status: bus.StatusDraining,
 	}))
 	di, err := svc.DescribeContainerInstances(context.Background(), &ecs.DescribeContainerInstancesInput{
@@ -415,13 +415,13 @@ func TestScheduler_ReapBucket_StopsStaleInstanceTasks(t *testing.T) {
 	taskID := containerInstanceShortID(aws.StringValue(out.Tasks[0].TaskArn))
 
 	// Backdate LastSeen beyond the heartbeat timeout.
-	kv, err := svc.bucket(testAccountID)
+	kv, err := svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
 	var rec InstanceRecord
-	_, err = getJSON(kv, InstanceKey("web", "i-1"), &rec)
+	_, err = getJSON(t.Context(), kv, InstanceKey("web", "i-1"), &rec)
 	require.NoError(t, err)
 	rec.LastSeen = time.Now().UTC().Add(-2 * heartbeatTimeout)
-	require.NoError(t, putJSON(kv, InstanceKey("web", "i-1"), &rec))
+	require.NoError(t, putJSON(t.Context(), kv, InstanceKey("web", "i-1"), &rec))
 
 	sc := NewScheduler(nc, svc, "test-holder")
 	sc.reapBucket(context.Background(), kv, testAccountID, time.Now().UTC())
@@ -444,9 +444,9 @@ func TestScheduler_AcquireLease_SingleLeader(t *testing.T) {
 	svc := NewService(nc, testRegion, "")
 	a := NewScheduler(nc, svc, "holder-a")
 	b := NewScheduler(nc, svc, "holder-b")
-	assert.True(t, a.acquireOrRefresh())
-	assert.False(t, b.acquireOrRefresh())
-	assert.True(t, a.acquireOrRefresh()) // refresh keeps leadership
+	assert.True(t, a.acquireOrRefresh(t.Context()))
+	assert.False(t, b.acquireOrRefresh(t.Context()))
+	assert.True(t, a.acquireOrRefresh(t.Context())) // refresh keeps leadership
 }
 
 // --- GPU placement dimension (Epic C2) ---
@@ -476,7 +476,7 @@ func registerTaskDefGPU(t *testing.T, svc *Service, family string, cpu, mem, gpu
 // capacity via the bus register path.
 func registerInstanceGPU(t *testing.T, svc *Service, cluster, id string, cpu, mem, gpu int) {
 	t.Helper()
-	require.NoError(t, svc.recordRegister(&bus.RegisterInstance{
+	require.NoError(t, svc.recordRegister(t.Context(), &bus.RegisterInstance{
 		AccountID:   testAccountID,
 		ClusterName: cluster,
 		InstanceID:  id,
@@ -526,10 +526,10 @@ func TestService_RunTask_GPU_ReservesOnPlacement(t *testing.T) {
 	assert.Empty(t, registeredGPU.StringSetValue)
 	assert.Empty(t, remainingGPU.StringSetValue)
 
-	kv, err := svc.bucket(testAccountID)
+	kv, err := svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
 	var rec InstanceRecord
-	found, err := getJSON(kv, InstanceKey("web", "i-1"), &rec)
+	found, err := getJSON(t.Context(), kv, InstanceKey("web", "i-1"), &rec)
 	require.NoError(t, err)
 	require.True(t, found)
 	assert.Equal(t, 2, rec.TotalGPU)
@@ -540,7 +540,7 @@ func TestService_RunTask_GPU_ReservesOnPlacement(t *testing.T) {
 		AccountID: testAccountID, ClusterName: "web", InstanceID: "i-1", TaskID: taskID,
 		LastStatus: bus.TaskStatusStopped, Reason: "exited",
 	}))
-	found, err = getJSON(kv, InstanceKey("web", "i-1"), &rec)
+	found, err = getJSON(t.Context(), kv, InstanceKey("web", "i-1"), &rec)
 	require.NoError(t, err)
 	require.True(t, found)
 	assert.Equal(t, 0, rec.ReservedGPU)
@@ -624,7 +624,7 @@ func TestRegisterContainerInstance_GPU_StringSetResource(t *testing.T) {
 // path (Epic C3 register report-back).
 func TestRecordRegister_GPU_CarriesDeviceUUIDs(t *testing.T) {
 	svc, _ := newTestService(t)
-	require.NoError(t, svc.recordRegister(&bus.RegisterInstance{
+	require.NoError(t, svc.recordRegister(t.Context(), &bus.RegisterInstance{
 		AccountID: testAccountID, ClusterName: "web", InstanceID: "i-1",
 		Capacity: bus.InstanceCapacity{CPU: 1024, MemoryMiB: 2048, GPU: 2, GPUIDs: []string{"GPU-aaa", "GPU-bbb"}},
 	}))

spinifex/handlers/ecs/services.go

@@ -16,6 +16,7 @@ import (
 	handlers_ec2_eip "github.com/mulgadc/spinifex/spinifex/handlers/ec2/eip"
 	handlers_elbv2 "github.com/mulgadc/spinifex/spinifex/handlers/elbv2"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // targetRegistrar registers/deregisters a service task's awsvpc ENI IP with an
@@ -136,12 +137,12 @@ func (s *Service) CreateService(ctx context.Context, input *ecs.CreateServiceInp
 	}
 
 	cluster := clusterShortName(aws.StringValue(input.Cluster))
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 	var clusterRec ClusterRecord
-	found, err := getJSON(kv, ClusterMetaKey(cluster), &clusterRec)
+	found, err := getJSON(ctx, kv, ClusterMetaKey(cluster), &clusterRec)
 	if err != nil {
 		return nil, err
 	}
@@ -150,13 +151,13 @@ func (s *Service) CreateService(ctx context.Context, input *ecs.CreateServiceInp
 	}
 
 	var existing ServiceRecord
-	if ok, gerr := getJSON(kv, ServiceKey(cluster, name), &existing); gerr != nil {
+	if ok, gerr := getJSON(ctx, kv, ServiceKey(cluster, name), &existing); gerr != nil {
 		return nil, gerr
 	} else if ok && existing.Status == ServiceStatusActive {
 		return &ecs.CreateServiceOutput{Service: s.serviceToAWS(accountID, &existing)}, nil
 	}
 
-	taskDef, err := s.resolveTaskDef(kv, aws.StringValue(input.TaskDefinition))
+	taskDef, err := s.resolveTaskDef(ctx, kv, aws.StringValue(input.TaskDefinition))
 	if err != nil {
 		return nil, err
 	}
@@ -190,7 +191,7 @@ func (s *Service) CreateService(ctx context.Context, input *ecs.CreateServiceInp
 	applyDeploymentConfig(&rec, input.DeploymentConfiguration)
 	rec.LastGoodTaskDefARN = taskDef.ARN
 	rec.Deployments = []Deployment{newPrimaryDeployment(rec.DeploymentID, taskDef, rec.DesiredCount)}
-	if err := putJSON(kv, ServiceKey(cluster, name), &rec); err != nil {
+	if err := putJSON(ctx, kv, ServiceKey(cluster, name), &rec); err != nil {
 		return nil, err
 	}
 
@@ -204,12 +205,12 @@ func (s *Service) CreateService(ctx context.Context, input *ecs.CreateServiceInp
 func (s *Service) UpdateService(ctx context.Context, input *ecs.UpdateServiceInput, accountID string) (*ecs.UpdateServiceOutput, error) {
 	cluster := clusterShortName(aws.StringValue(input.Cluster))
 	name := serviceShortName(aws.StringValue(input.Service))
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 	var rec ServiceRecord
-	found, err := getJSON(kv, ServiceKey(cluster, name), &rec)
+	found, err := getJSON(ctx, kv, ServiceKey(cluster, name), &rec)
 	if err != nil {
 		return nil, err
 	}
@@ -225,7 +226,7 @@ func (s *Service) UpdateService(ctx context.Context, input *ecs.UpdateServiceInp
 		rec.DesiredCount = int(aws.Int64Value(input.DesiredCount))
 	}
 	if td := aws.StringValue(input.TaskDefinition); td != "" {
-		taskDef, terr := s.resolveTaskDef(kv, td)
+		taskDef, terr := s.resolveTaskDef(ctx, kv, td)
 		if terr != nil {
 			return nil, terr
 		}
@@ -236,7 +237,7 @@ func (s *Service) UpdateService(ctx context.Context, input *ecs.UpdateServiceInp
 		}
 	}
 	rec.UpdatedAt = time.Now().UTC()
-	if err := putJSON(kv, ServiceKey(cluster, name), &rec); err != nil {
+	if err := putJSON(ctx, kv, ServiceKey(cluster, name), &rec); err != nil {
 		return nil, err
 	}
 	if err := s.reconcileService(ctx, kv, accountID, &rec); err != nil {
@@ -250,12 +251,12 @@ func (s *Service) UpdateService(ctx context.Context, input *ecs.UpdateServiceInp
 func (s *Service) DeleteService(ctx context.Context, input *ecs.DeleteServiceInput, accountID string) (*ecs.DeleteServiceOutput, error) {
 	cluster := clusterShortName(aws.StringValue(input.Cluster))
 	name := serviceShortName(aws.StringValue(input.Service))
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 	var rec ServiceRecord
-	found, err := getJSON(kv, ServiceKey(cluster, name), &rec)
+	found, err := getJSON(ctx, kv, ServiceKey(cluster, name), &rec)
 	if err != nil {
 		return nil, err
 	}
@@ -264,7 +265,7 @@ func (s *Service) DeleteService(ctx context.Context, input *ecs.DeleteServiceInp
 	}
 
 	rec.DesiredCount = 0
-	tasks, err := s.listServiceTasks(kv, cluster, name)
+	tasks, err := s.listServiceTasks(ctx, kv, cluster, name)
 	if err != nil {
 		return nil, err
 	}
@@ -275,16 +276,16 @@ func (s *Service) DeleteService(ctx context.Context, input *ecs.DeleteServiceInp
 	rec.RunningCount = 0
 	rec.PendingCount = 0
 	rec.UpdatedAt = time.Now().UTC()
-	if err := putJSON(kv, ServiceKey(cluster, name), &rec); err != nil {
+	if err := putJSON(ctx, kv, ServiceKey(cluster, name), &rec); err != nil {
 		return nil, err
 	}
 	return &ecs.DeleteServiceOutput{Service: s.serviceToAWS(accountID, &rec)}, nil
 }
 
 // DescribeServices returns the named services; unknown names surface as failures.
-func (s *Service) DescribeServices(_ context.Context, input *ecs.DescribeServicesInput, accountID string) (*ecs.DescribeServicesOutput, error) {
+func (s *Service) DescribeServices(ctx context.Context, input *ecs.DescribeServicesInput, accountID string) (*ecs.DescribeServicesOutput, error) {
 	cluster := clusterShortName(aws.StringValue(input.Cluster))
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
@@ -292,7 +293,7 @@ func (s *Service) DescribeServices(_ context.Context, input *ecs.DescribeService
 	for _, ref := range awsStringSlice(input.Services) {
 		name := serviceShortName(ref)
 		var rec ServiceRecord
-		found, gerr := getJSON(kv, ServiceKey(cluster, name), &rec)
+		found, gerr := getJSON(ctx, kv, ServiceKey(cluster, name), &rec)
 		if gerr != nil {
 			return nil, gerr
 		}
@@ -306,13 +307,13 @@ func (s *Service) DescribeServices(_ context.Context, input *ecs.DescribeService
 }
 
 // ListServices returns the ARNs of every service in a cluster.
-func (s *Service) ListServices(_ context.Context, input *ecs.ListServicesInput, accountID string) (*ecs.ListServicesOutput, error) {
+func (s *Service) ListServices(ctx context.Context, input *ecs.ListServicesInput, accountID string) (*ecs.ListServicesOutput, error) {
 	cluster := clusterShortName(aws.StringValue(input.Cluster))
-	kv, err := s.bucket(accountID)
+	kv, err := s.bucket(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
-	recs, err := s.listServiceRecords(kv, cluster)
+	recs, err := s.listServiceRecords(ctx, kv, cluster)
 	if err != nil {
 		return nil, err
 	}
@@ -330,14 +331,14 @@ func (s *Service) ListServices(_ context.Context, input *ecs.ListServicesInput,
 // and drains superseded (old) tasks while healthy running stays at/above
 // minimumHealthyPercent, then advances the rollout / circuit-breaker state. Per-
 // deployment and overall counts are refreshed as a side effect.
-func (s *Service) reconcileService(ctx context.Context, kv nats.KeyValue, accountID string, svc *ServiceRecord) error {
+func (s *Service) reconcileService(ctx context.Context, kv jetstream.KeyValue, accountID string, svc *ServiceRecord) error {
 	if svc.Status != ServiceStatusActive {
 		return nil
 	}
 	svc.normalizeDeploymentConfig()
 	svc.ensurePrimaryDeployment()
 
-	tasks, err := s.listServiceTasks(kv, svc.Cluster, svc.Name)
+	tasks, err := s.listServiceTasks(ctx, kv, svc.Cluster, svc.Name)
 	if err != nil {
 		return err
 	}
@@ -349,7 +350,7 @@ func (s *Service) reconcileService(ctx context.Context, kv nats.KeyValue, accoun
 		// let the next tick roll the replacement in.
 		if tripCircuitBreaker(svc, primary) {
 			svc.RunningCount, svc.PendingCount = running, pending
-			return putJSON(kv, ServiceKey(svc.Cluster, svc.Name), svc)
+			return putJSON(ctx, kv, ServiceKey(svc.Cluster, svc.Name), svc)
 		}
 		desired := svc.DesiredCount
 		maxCount := desired * svc.MaximumPercent / 100
@@ -365,43 +366,48 @@ func (s *Service) reconcileService(ctx context.Context, kv nats.KeyValue, accoun
 		s.stopSurplusTasks(ctx, kv, accountID, tasks, primary.ID, desired, running, minCount)
 
 		// Re-tally after launch/stop so rollout state + persisted counts are current.
-		if fresh, ferr := s.listServiceTasks(kv, svc.Cluster, svc.Name); ferr == nil {
+		if fresh, ferr := s.listServiceTasks(ctx, kv, svc.Cluster, svc.Name); ferr == nil {
 			running, pending = tallyDeployments(svc, fresh)
 		}
 		updateRolloutState(svc, primary, desired)
 	}
 
 	svc.RunningCount, svc.PendingCount = running, pending
-	return putJSON(kv, ServiceKey(svc.Cluster, svc.Name), svc)
+	return putJSON(ctx, kv, ServiceKey(svc.Cluster, svc.Name), svc)
 }
 
 // reconcileAllServices is the scheduler-tick fan-out: every ACTIVE service in
 // every ECS account bucket is reconciled. Runs only on the scheduler leader.
-func (s *Service) reconcileAllServices() {
-	ctx := context.Background()
+// Returns an error when the account enumeration could not be completed, so a
+// pass that could not see the whole fleet is reported rather than read as "no
+// services to reconcile" — every unlisted account stalls below its desired count.
+func (s *Service) reconcileAllServices(ctx context.Context) error {
 	js, err := s.js()
 	if err != nil {
-		return
+		return err
 	}
-	for bucket := range js.KeyValueStoreNames() {
-		accountID, ok := accountIDFromBucket(bucket)
-		if !ok {
+	buckets, err := accountBuckets(ctx, s.nc)
+	if err != nil {
+		return err
+	}
+	for _, bucket := range buckets {
+		kv, err := js.KeyValue(ctx, bucket.name)
+		if err != nil {
+			slog.Error("ECS reconcile: open bucket failed", "bucket", bucket.name, "err", err)
 			continue
 		}
-		kv, err := js.KeyValue(bucket)
-		if err != nil {
-			continue
-		}
-		recs, err := s.allServiceRecords(kv)
+		recs, err := s.allServiceRecords(ctx, kv)
 		if err != nil {
+			slog.Error("ECS reconcile: list services failed", "bucket", bucket.name, "err", err)
 			continue
 		}
 		for i := range recs {
-			if err := s.reconcileService(ctx, kv, accountID, &recs[i]); err != nil {
+			if err := s.reconcileService(ctx, kv, bucket.accountID, &recs[i]); err != nil {
 				slog.Error("ECS reconcile: service failed", "service", recs[i].Name, "err", err)
 			}
 		}
 	}
+	return nil
 }
 
 // launchDeploymentTasks places n tasks for a specific deployment via the standard
@@ -443,7 +449,7 @@ func (s *Service) launchDeploymentTasks(ctx context.Context, accountID string, s
 // deployment tasks toward zero, and the primary's own surplus on a scale-in.
 // PENDING surplus is stopped freely; RUNNING surplus is gated so healthy running
 // never drops below minimumHealthyPercent (minCount).
-func (s *Service) stopSurplusTasks(ctx context.Context, kv nats.KeyValue, accountID string, tasks []TaskRecord, primaryID string, desired, runningTotal, minCount int) {
+func (s *Service) stopSurplusTasks(ctx context.Context, kv jetstream.KeyValue, accountID string, tasks []TaskRecord, primaryID string, desired, runningTotal, minCount int) {
 	var oldPending, oldRunning, primaryPending, primaryRunning []int
 	for i := range tasks {
 		isPrimary := deploymentIDFromStartedBy(tasks[i].StartedBy) == primaryID
@@ -501,9 +507,9 @@ func (s *Service) stopSurplusTasks(ctx context.Context, kv nats.KeyValue, accoun
 // STOPPED nor already requested to stop (DesiredStatus=STOPPED). A cooperatively
 // stopped task lingers RUNNING until the agent reaps it, but it is on its way out,
 // so it must not count toward the service's running total for scaling decisions.
-func (s *Service) listServiceTasks(kv nats.KeyValue, cluster, name string) ([]TaskRecord, error) {
+func (s *Service) listServiceTasks(ctx context.Context, kv jetstream.KeyValue, cluster, name string) ([]TaskRecord, error) {
 	group := serviceTaskGroup(name)
-	all, err := s.listTaskRecords(kv, cluster)
+	all, err := s.listTaskRecords(ctx, kv, cluster)
 	if err != nil {
 		return nil, err
 	}
@@ -517,15 +523,15 @@ func (s *Service) listServiceTasks(kv nats.KeyValue, cluster, name string) ([]Ta
 }
 
 // listTaskRecords returns every task record in a cluster.
-func (s *Service) listTaskRecords(kv nats.KeyValue, cluster string) ([]TaskRecord, error) {
-	keys, err := keysWithPrefix(kv, TasksPrefix(cluster))
+func (s *Service) listTaskRecords(ctx context.Context, kv jetstream.KeyValue, cluster string) ([]TaskRecord, error) {
+	keys, err := keysWithPrefix(ctx, kv, TasksPrefix(cluster))
 	if err != nil {
 		return nil, err
 	}
 	out := make([]TaskRecord, 0, len(keys))
 	for _, k := range keys {
 		var rec TaskRecord
-		found, err := getJSON(kv, k, &rec)
+		found, err := getJSON(ctx, kv, k, &rec)
 		if err != nil {
 			return nil, err
 		}
@@ -537,15 +543,15 @@ func (s *Service) listTaskRecords(kv nats.KeyValue, cluster string) ([]TaskRecor
 }
 
 // listServiceRecords returns every service record in a cluster.
-func (s *Service) listServiceRecords(kv nats.KeyValue, cluster string) ([]ServiceRecord, error) {
-	keys, err := keysWithPrefix(kv, ServicesPrefix(cluster))
+func (s *Service) listServiceRecords(ctx context.Context, kv jetstream.KeyValue, cluster string) ([]ServiceRecord, error) {
+	keys, err := keysWithPrefix(ctx, kv, ServicesPrefix(cluster))
 	if err != nil {
 		return nil, err
 	}
 	out := make([]ServiceRecord, 0, len(keys))
 	for _, k := range keys {
 		var rec ServiceRecord
-		found, err := getJSON(kv, k, &rec)
+		found, err := getJSON(ctx, kv, k, &rec)
 		if err != nil {
 			return nil, err
 		}
@@ -557,8 +563,8 @@ func (s *Service) listServiceRecords(kv nats.KeyValue, cluster string) ([]Servic
 }
 
 // allServiceRecords returns every service record across all clusters in a bucket.
-func (s *Service) allServiceRecords(kv nats.KeyValue) ([]ServiceRecord, error) {
-	keys, err := keysWithPrefix(kv, "clusters/")
+func (s *Service) allServiceRecords(ctx context.Context, kv jetstream.KeyValue) ([]ServiceRecord, error) {
+	keys, err := keysWithPrefix(ctx, kv, "clusters/")
 	if err != nil {
 		return nil, err
 	}
@@ -568,7 +574,7 @@ func (s *Service) allServiceRecords(kv nats.KeyValue) ([]ServiceRecord, error) {
 			continue
 		}
 		var rec ServiceRecord
-		found, err := getJSON(kv, k, &rec)
+		found, err := getJSON(ctx, kv, k, &rec)
 		if err != nil {
 			return nil, err
 		}

spinifex/handlers/ecs/services_test.go

@@ -8,7 +8,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ecs"
 	"github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus"
-	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -34,7 +34,7 @@ func (r *stubRegistrar) Deregister(_ context.Context, _, tg, ip string, port int
 
 // serviceTestRig wires a Service with a stub registrar and a seeded cluster +
 // task definition + one fat container instance.
-func serviceTestRig(t *testing.T) (*Service, *stubRegistrar, nats.KeyValue) {
+func serviceTestRig(t *testing.T) (*Service, *stubRegistrar, jetstream.KeyValue) {
 	t.Helper()
 	svc, _ := newTestService(t)
 	reg := &stubRegistrar{}
@@ -43,7 +43,7 @@ func serviceTestRig(t *testing.T) (*Service, *stubRegistrar, nats.KeyValue) {
 	require.NoError(t, err)
 	registerTaskDef(t, svc, "app", 128, 256)
 	registerInstance(t, svc, "web", "i-1", 4096, 8192)
-	kv, err := svc.bucket(testAccountID)
+	kv, err := svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
 	return svc, reg, kv
 }
@@ -62,7 +62,7 @@ func TestService_CreateService_LaunchesReplicas(t *testing.T) {
 	assert.Equal(t, ServiceARN(testRegion, testAccountID, "web", "web"), aws.StringValue(out.Service.ServiceArn))
 	assert.Equal(t, int64(2), aws.Int64Value(out.Service.PendingCount))
 
-	tasks, err := svc.listServiceTasks(kv, "web", "web")
+	tasks, err := svc.listServiceTasks(t.Context(), kv, "web", "web")
 	require.NoError(t, err)
 	require.Len(t, tasks, 2)
 	for _, task := range tasks {
@@ -106,7 +106,7 @@ func TestService_Reconcile_ReplacesStoppedTask(t *testing.T) {
 	}, testAccountID)
 	require.NoError(t, err)
 
-	tasks, err := svc.listServiceTasks(kv, "web", "web")
+	tasks, err := svc.listServiceTasks(t.Context(), kv, "web", "web")
 	require.NoError(t, err)
 	require.Len(t, tasks, 2)
 
@@ -115,13 +115,13 @@ func TestService_Reconcile_ReplacesStoppedTask(t *testing.T) {
 		AccountID: testAccountID, ClusterName: "web", InstanceID: "i-1",
 		TaskID: tasks[0].TaskID, LastStatus: bus.TaskStatusStopped,
 	}))
-	live, err := svc.listServiceTasks(kv, "web", "web")
+	live, err := svc.listServiceTasks(t.Context(), kv, "web", "web")
 	require.NoError(t, err)
 	require.Len(t, live, 1) // STOPPED task drops out of the active set
 
 	// Reconcile launches a replacement back to desired=2.
-	svc.reconcileAllServices()
-	live, err = svc.listServiceTasks(kv, "web", "web")
+	require.NoError(t, svc.reconcileAllServices(t.Context()))
+	live, err = svc.listServiceTasks(t.Context(), kv, "web", "web")
 	require.NoError(t, err)
 	assert.Len(t, live, 2)
 }
@@ -134,7 +134,7 @@ func TestService_UpdateService_ScalesIn(t *testing.T) {
 		DesiredCount: aws.Int64(3),
 	}, testAccountID)
 	require.NoError(t, err)
-	live, _ := svc.listServiceTasks(kv, "web", "web")
+	live, _ := svc.listServiceTasks(t.Context(), kv, "web", "web")
 	require.Len(t, live, 3)
 
 	out, err := svc.UpdateService(context.Background(), &ecs.UpdateServiceInput{
@@ -143,7 +143,7 @@ func TestService_UpdateService_ScalesIn(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, int64(1), aws.Int64Value(out.Service.DesiredCount))
 
-	live, err = svc.listServiceTasks(kv, "web", "web")
+	live, err = svc.listServiceTasks(t.Context(), kv, "web", "web")
 	require.NoError(t, err)
 	assert.Len(t, live, 1)
 }
@@ -162,7 +162,7 @@ func TestService_DeleteService_DrainsAndInactivates(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, ServiceStatusInactive, aws.StringValue(out.Service.Status))
 
-	live, err := svc.listServiceTasks(kv, "web", "web")
+	live, err := svc.listServiceTasks(t.Context(), kv, "web", "web")
 	require.NoError(t, err)
 	assert.Empty(t, live)
 
@@ -172,8 +172,8 @@ func TestService_DeleteService_DrainsAndInactivates(t *testing.T) {
 	}, testAccountID)
 	require.NoError(t, err)
 	require.Len(t, desc.Services, 1)
-	svc.reconcileAllServices()
-	live, _ = svc.listServiceTasks(kv, "web", "web")
+	require.NoError(t, svc.reconcileAllServices(t.Context()))
+	live, _ = svc.listServiceTasks(t.Context(), kv, "web", "web")
 	assert.Empty(t, live)
 }
 
@@ -214,7 +214,7 @@ func TestService_StartTask_PlacesPerInstance(t *testing.T) {
 	assert.Empty(t, out.Failures)
 	assert.Equal(t, TaskStatusPending, aws.StringValue(out.Tasks[0].LastStatus))
 
-	tasks, err := svc.listTaskRecords(kv, "web")
+	tasks, err := svc.listTaskRecords(t.Context(), kv, "web")
 	require.NoError(t, err)
 	require.Len(t, tasks, 1)
 	assert.Equal(t, "custom", tasks[0].Group)
@@ -263,7 +263,7 @@ func TestService_StopTask_StopsAndDeregisters(t *testing.T) {
 		NetworkMode: NetworkModeAwsvpc, ENIPrivateIP: "10.0.1.9",
 		ReservedCPU: 64, ReservedMemoryMiB: 64,
 	}
-	require.NoError(t, putJSON(kv, TaskKey("web", "t-9"), task))
+	require.NoError(t, putJSON(t.Context(), kv, TaskKey("web", "t-9"), task))
 
 	// Cooperative stop: the task stays RUNNING with desiredStatus=STOPPED and a
 	// stop directive lands in the instance inbox; nothing is deregistered until the
@@ -277,7 +277,7 @@ func TestService_StopTask_StopsAndDeregisters(t *testing.T) {
 	assert.Empty(t, reg.deregistered)
 
 	var sd bus.StopDirective
-	found, err := getJSON(kv, StopKey("web", "i-1", "t-9"), &sd)
+	found, err := getJSON(t.Context(), kv, StopKey("web", "i-1", "t-9"), &sd)
 	require.NoError(t, err)
 	require.True(t, found)
 	assert.Equal(t, "bye", sd.Reason)
@@ -288,7 +288,7 @@ func TestService_StopTask_StopsAndDeregisters(t *testing.T) {
 		TaskID: "t-9", LastStatus: bus.TaskStatusStopped,
 	}))
 	assert.Equal(t, []string{key(tgARN, "10.0.1.9", 80)}, reg.deregistered)
-	found, err = getJSON(kv, StopKey("web", "i-1", "t-9"), &sd)
+	found, err = getJSON(t.Context(), kv, StopKey("web", "i-1", "t-9"), &sd)
 	require.NoError(t, err)
 	assert.False(t, found)
 
@@ -312,7 +312,7 @@ func TestService_StopTask_UnplacedForceStops(t *testing.T) {
 		TaskID: "t-np", Cluster: "web",
 		LastStatus: TaskStatusPending, DesiredStatus: TaskStatusRunning,
 	}
-	require.NoError(t, putJSON(kv, TaskKey("web", "t-np"), task))
+	require.NoError(t, putJSON(t.Context(), kv, TaskKey("web", "t-np"), task))
 
 	out, err := svc.StopTask(context.Background(), &ecs.StopTaskInput{
 		Cluster: aws.String("web"), Task: aws.String("t-np"), Reason: aws.String("cancel"),
@@ -344,7 +344,7 @@ func TestService_TargetGroup_RegisterOnRunningDeregisterOnStopped(t *testing.T)
 		ContainerInstanceID: "i-1", LastStatus: TaskStatusPending,
 		NetworkMode: NetworkModeAwsvpc, ENIPrivateIP: "10.0.1.5",
 	}
-	require.NoError(t, putJSON(kv, TaskKey("web", "t-1"), task))
+	require.NoError(t, putJSON(t.Context(), kv, TaskKey("web", "t-1"), task))
 
 	// RUNNING → register.
 	require.NoError(t, svc.recordTaskState(context.Background(), &bus.TaskState{

spinifex/handlers/ecs/store.go

@@ -1,14 +1,16 @@
 package handlers_ecs
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"log/slog"
 	"time"
 
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
-	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // JetStream KV bucket and key-path constants for the ECS control plane.
@@ -169,16 +171,46 @@ func AccountBucketName(accountID string) string {
 	return KVBucketECSAccountPrefix + accountID
 }
 
+// accountBucket pairs a per-account KV bucket name with the account it belongs
+// to, so the sweeps that walk every account do not each re-derive the ID.
+type accountBucket struct {
+	name      string
+	accountID string
+}
+
+// accountBuckets returns every ECS per-account bucket. It fails rather than
+// returning a short list when the enumeration could not be completed: the lister
+// behind it closes its channel identically on success and on error, so a sweep
+// that ignored the failure would read an unreachable JetStream as an empty fleet
+// and silently skip every account it could not see.
+func accountBuckets(ctx context.Context, nc *nats.Conn) ([]accountBucket, error) {
+	js, err := jetstream.New(nc)
+	if err != nil {
+		return nil, fmt.Errorf("jetstream: %w", err)
+	}
+	names, err := kvutil.BucketNames(ctx, js)
+	if err != nil {
+		return nil, err
+	}
+	buckets := make([]accountBucket, 0, len(names))
+	for _, name := range names {
+		if accountID, ok := accountIDFromBucket(name); ok {
+			buckets = append(buckets, accountBucket{name: name, accountID: accountID})
+		}
+	}
+	return buckets, nil
+}
+
 // GetOrCreateAccountBucket returns the per-account KV bucket for accountID,
 // creating it on first use. Idempotent: subsequent calls with the same accountID
 // return the existing handle.
-func GetOrCreateAccountBucket(js nats.JetStreamContext, accountID string) (nats.KeyValue, error) {
+func GetOrCreateAccountBucket(ctx context.Context, js jetstream.JetStream, accountID string) (jetstream.KeyValue, error) {
 	bucket := AccountBucketName(accountID)
-	kv, err := utils.GetOrCreateKVBucket(js, bucket, KVBucketECSAccountHistory)
+	kv, err := kvutil.GetOrCreateBucket(ctx, js, bucket, KVBucketECSAccountHistory)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create ECS per-account KV bucket %s: %w", bucket, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(bucket, kv, KVBucketECSAccountVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, bucket, kv, KVBucketECSAccountVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", bucket, err)
 	}
 	return kv, nil
@@ -187,22 +219,22 @@ func GetOrCreateAccountBucket(js nats.JetStreamContext, accountID string) (nats.
 // InitLeaderBucket creates (or attaches to) the shared spinifex-ecs-leader bucket
 // used for per-cluster scheduler leader-lease CAS locks. The bucket is configured
 // with History=1 and a 60s TTL so stale leases expire on their own when a leader
-// dies mid-cycle. utils.GetOrCreateKVBucket doesn't expose a TTL knob, so this
+// dies mid-cycle. kvutil.GetOrCreateBucket doesn't expose a TTL knob, so this
 // takes the direct js.CreateKeyValue path and falls back to js.KeyValue on
 // already-exists.
-func InitLeaderBucket(js nats.JetStreamContext) (nats.KeyValue, error) {
-	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{
+func InitLeaderBucket(ctx context.Context, js jetstream.JetStream) (jetstream.KeyValue, error) {
+	kv, err := js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
 		Bucket:  KVBucketECSLeader,
 		History: 1,
 		TTL:     KVBucketECSLeaderTTL,
 	})
 	if err != nil {
-		kv, err = js.KeyValue(KVBucketECSLeader)
+		kv, err = js.KeyValue(ctx, KVBucketECSLeader)
 		if err != nil {
 			return nil, fmt.Errorf("failed to create or open ECS leader bucket %s: %w", KVBucketECSLeader, err)
 		}
 	}
-	if err := migrate.DefaultRegistry.RunKV(KVBucketECSLeader, kv, KVBucketECSLeaderVersion); err != nil {
+	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketECSLeader, kv, KVBucketECSLeaderVersion); err != nil {
 		return nil, fmt.Errorf("migrate %s: %w", KVBucketECSLeader, err)
 	}
 	slog.Info("ECS leader bucket initialized", "bucket", KVBucketECSLeader, "ttl", KVBucketECSLeaderTTL)