feat: jetstream kv for instance state

cmd/hive/cmd/templates/nats.conf

@@ -32,12 +32,12 @@ authorization {
 #  token: "{{.NatsToken}}"
 }
 
-# JetStream for persistent messaging
-#jetstream {
-#  store_dir: "{{.DataDir}}/nats/"
-#  max_memory_store: 1GB
-#  max_file_store: 10GB
-#}
+# JetStream for persistent messaging and KV store
+jetstream {
+  store_dir: "{{.DataDir}}/nats/"
+  max_memory_store: 1GB
+  max_file_store: 10GB
+}
 
 # Logging
 debug: false

go.mod

@@ -49,6 +49,7 @@ require (
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/klauspost/compress v1.18.0 // indirect
+	github.com/klauspost/cpuid/v2 v2.3.0 // indirect
 	github.com/lithammer/fuzzysearch v1.1.8 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect

go.sum

@@ -103,8 +103,8 @@ github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYW
 github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
 github.com/klauspost/cpuid/v2 v2.0.10/go.mod h1:g2LTdtYhdyuGPqyWyv7qRAmj1WBqxuObKfj5c0PQa7c=
 github.com/klauspost/cpuid/v2 v2.0.12/go.mod h1:g2LTdtYhdyuGPqyWyv7qRAmj1WBqxuObKfj5c0PQa7c=
-github.com/klauspost/cpuid/v2 v2.2.3 h1:sxCkb+qR91z4vsqw4vGGZlDgPz3G7gjaLyK3V8y70BU=
-github.com/klauspost/cpuid/v2 v2.2.3/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
+github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y=
+github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=

hive/daemon/daemon.go

@@ -93,6 +93,9 @@ type Daemon struct {
 	startTime  time.Time
 	configPath string
 
+	// JetStream manager for KV state storage (nil if JetStream disabled)
+	jsManager *JetStreamManager
+
 	mu sync.Mutex
 }
 
@@ -405,8 +408,19 @@ func (d *Daemon) Start() error {
 
 	log.Printf("Connected to NATS server at %s", d.config.NATS.Host)
 
-	// Load existing state for VMs
-	// Load state from disk
+	// Initialize JetStream for KV state storage (falls back to disk if unavailable)
+	d.jsManager, err = NewJetStreamManager(d.natsConn)
+	if err != nil {
+		slog.Warn("Failed to init JetStream, falling back to file", "error", err)
+		d.jsManager = nil
+	} else if err := d.jsManager.InitKVBucket(); err != nil {
+		slog.Warn("Failed to init KV bucket, falling back to file", "error", err)
+		d.jsManager = nil
+	} else {
+		slog.Info("JetStream KV store initialized successfully")
+	}
+
+	// Load existing state for VMs from JetStream or disk
 	err = d.LoadState()
 	if err != nil {
 		slog.Warn("Failed to load state from disk, continuing with empty state", "error", err)
@@ -773,8 +787,20 @@ func (d *Daemon) ClusterManager() error {
 	return nil
 }
 
-// Write the state to disk
+// WriteState writes the instance state to JetStream KV store (if enabled) with disk fallback
 func (d *Daemon) WriteState() error {
+	if d.jsManager != nil {
+		if err := d.jsManager.WriteState(d.node, &d.Instances); err != nil {
+			slog.Error("JetStream write failed, falling back to disk", "error", err)
+			return d.writeStateToDisk()
+		}
+		return nil
+	}
+	return d.writeStateToDisk()
+}
+
+// writeStateToDisk writes the instance state to disk (original WriteState behavior)
+func (d *Daemon) writeStateToDisk() error {
 	d.Instances.Mu.Lock()
 	defer d.Instances.Mu.Unlock()
 
@@ -828,8 +854,37 @@ func (d *Daemon) InitaliseVMs() {
 	// Step 2: Loop through each instance and start it
 }
 
-// Load state from disk
+// LoadState loads the instance state from JetStream KV store (if enabled) with disk fallback
 func (d *Daemon) LoadState() error {
+	if d.jsManager != nil {
+		instances, err := d.jsManager.LoadState(d.node)
+		if err != nil {
+			slog.Warn("JetStream load failed, trying disk", "error", err)
+			return d.loadStateFromDisk()
+		}
+
+		// If JetStream returned empty state, check if disk has data to migrate
+		if len(instances.VMS) == 0 {
+			diskErr := d.loadStateFromDisk()
+			if diskErr == nil && len(d.Instances.VMS) > 0 {
+				slog.Info("Migrating state from disk to JetStream", "instances", len(d.Instances.VMS))
+				// Write existing disk state to JetStream for future use
+				if writeErr := d.jsManager.WriteState(d.node, &d.Instances); writeErr != nil {
+					slog.Warn("Failed to migrate state to JetStream", "error", writeErr)
+				}
+				return nil // Use the disk state we just loaded
+			}
+		}
+
+		// Copy only the VMS map, not the mutex
+		d.Instances.VMS = instances.VMS
+		return nil
+	}
+	return d.loadStateFromDisk()
+}
+
+// loadStateFromDisk loads the instance state from disk (original LoadState behavior)
+func (d *Daemon) loadStateFromDisk() error {
 	configPath := fmt.Sprintf("%s/%s", d.config.BaseDir, "instances.json")
 	jsonData, err := os.ReadFile(configPath)
 	if err != nil {
@@ -1420,6 +1475,12 @@ func (d *Daemon) setupShutdown() {
 
 		}
 
+		// Write state to JetStream before closing NATS connection
+		err := d.WriteState()
+		if err != nil {
+			slog.Error("Failed to write state", "err", err)
+		}
+
 		// Close NATS connection
 		d.natsConn.Close()
 
@@ -1431,12 +1492,6 @@ func (d *Daemon) setupShutdown() {
 			}
 		}
 
-		// Write the state to disk
-		err := d.WriteState()
-		if err != nil {
-			slog.Error("Failed to write state to disk", "err", err)
-		}
-
 		// Wait for any ongoing operations to complete
 		// TODO: Implement cleanup of running instances
 		log.Println("Shutdown complete")

hive/daemon/jetstream.go

@@ -0,0 +1,140 @@
+package daemon
+
+import (
+	"encoding/json"
+	"errors"
+	"log/slog"
+
+	"github.com/mulgadc/hive/hive/vm"
+	"github.com/nats-io/nats.go"
+)
+
+const (
+	// InstanceStateBucket is the name of the KV bucket for storing instance state
+	InstanceStateBucket = "hive-instance-state"
+	// InstanceStatePrefix is the key prefix for instance state entries
+	InstanceStatePrefix = "node."
+)
+
+// JetStreamManager manages JetStream KV store operations for instance state
+type JetStreamManager struct {
+	js nats.JetStreamContext
+	kv nats.KeyValue
+}
+
+// NewJetStreamManager creates a new JetStreamManager from a NATS connection
+func NewJetStreamManager(nc *nats.Conn) (*JetStreamManager, error) {
+	js, err := nc.JetStream()
+	if err != nil {
+		return nil, err
+	}
+
+	return &JetStreamManager{
+		js: js,
+	}, nil
+}
+
+// InitKVBucket initializes the KV bucket, creating it if it doesn't exist
+func (m *JetStreamManager) InitKVBucket() error {
+	// Try to get the existing bucket first
+	kv, err := m.js.KeyValue(InstanceStateBucket)
+	if err != nil {
+		if errors.Is(err, nats.ErrBucketNotFound) {
+			// Bucket doesn't exist, create it
+			slog.Info("Creating JetStream KV bucket", "bucket", InstanceStateBucket)
+			kv, err = m.js.CreateKeyValue(&nats.KeyValueConfig{
+				Bucket:      InstanceStateBucket,
+				Description: "Hive instance state storage",
+				History:     1, // Only keep latest value
+			})
+			if err != nil {
+				return err
+			}
+		} else {
+			return err
+		}
+	} else {
+		slog.Info("Connected to existing JetStream KV bucket", "bucket", InstanceStateBucket)
+	}
+
+	m.kv = kv
+	return nil
+}
+
+// WriteState writes the instance state to the KV store for the given node
+func (m *JetStreamManager) WriteState(nodeID string, instances *vm.Instances) error {
+	if m.kv == nil {
+		return errors.New("KV bucket not initialized")
+	}
+
+	instances.Mu.Lock()
+	defer instances.Mu.Unlock()
+
+	// Create a struct without the mutex to avoid copying the lock
+	state := struct {
+		VMS map[string]*vm.VM `json:"vms"`
+	}{
+		VMS: instances.VMS,
+	}
+
+	jsonData, err := json.Marshal(state)
+	if err != nil {
+		return err
+	}
+
+	key := InstanceStatePrefix + nodeID
+	_, err = m.kv.Put(key, jsonData)
+	if err != nil {
+		return err
+	}
+
+	slog.Debug("Wrote state to JetStream KV", "key", key, "instances", len(instances.VMS))
+	return nil
+}
+
+// LoadState loads the instance state from the KV store for the given node
+func (m *JetStreamManager) LoadState(nodeID string) (*vm.Instances, error) {
+	if m.kv == nil {
+		return nil, errors.New("KV bucket not initialized")
+	}
+
+	key := InstanceStatePrefix + nodeID
+	entry, err := m.kv.Get(key)
+	if err != nil {
+		if errors.Is(err, nats.ErrKeyNotFound) {
+			// Key doesn't exist, return empty state
+			slog.Debug("No existing state in JetStream KV, returning empty state", "key", key)
+			return &vm.Instances{VMS: make(map[string]*vm.VM)}, nil
+		}
+		return nil, err
+	}
+
+	var instances vm.Instances
+	if err := json.Unmarshal(entry.Value(), &instances); err != nil {
+		return nil, err
+	}
+
+	// Ensure the VMS map is initialized
+	if instances.VMS == nil {
+		instances.VMS = make(map[string]*vm.VM)
+	}
+
+	slog.Debug("Loaded state from JetStream KV", "key", key, "instances", len(instances.VMS))
+	return &instances, nil
+}
+
+// DeleteState removes the instance state from the KV store for the given node
+func (m *JetStreamManager) DeleteState(nodeID string) error {
+	if m.kv == nil {
+		return errors.New("KV bucket not initialized")
+	}
+
+	key := InstanceStatePrefix + nodeID
+	err := m.kv.Delete(key)
+	if err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
+		return err
+	}
+
+	slog.Debug("Deleted state from JetStream KV", "key", key)
+	return nil
+}

hive/daemon/jetstream_test.go

@@ -0,0 +1,386 @@
+package daemon
+
+import (
+	"os"
+	"testing"
+	"time"
+
+	"github.com/mulgadc/hive/hive/vm"
+	"github.com/nats-io/nats-server/v2/server"
+	"github.com/nats-io/nats.go"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// startTestNATSServerWithJetStream starts an embedded NATS server with JetStream enabled for testing
+func startTestNATSServerWithJetStream(t *testing.T) (*server.Server, string) {
+	// Create a temporary directory for JetStream storage
+	tmpDir, err := os.MkdirTemp("", "nats-js-test-*")
+	require.NoError(t, err, "Failed to create temp directory for JetStream")
+
+	t.Cleanup(func() {
+		os.RemoveAll(tmpDir)
+	})
+
+	opts := &server.Options{
+		Host:      "127.0.0.1",
+		Port:      -1, // Let NATS auto-allocate an available port
+		JetStream: true,
+		StoreDir:  tmpDir,
+		NoLog:     true,
+		NoSigs:    true,
+	}
+
+	ns, err := server.NewServer(opts)
+	require.NoError(t, err, "Failed to create NATS server with JetStream")
+
+	// Start server in a goroutine
+	go ns.Start()
+
+	// Wait for server to be ready
+	if !ns.ReadyForConnections(5 * time.Second) {
+		t.Fatal("NATS server with JetStream failed to start")
+	}
+
+	// Get the actual URL that was assigned
+	url := ns.ClientURL()
+	t.Logf("Test NATS server with JetStream started at: %s", url)
+
+	return ns, url
+}
+
+// TestJetStreamManager_WriteAndLoadState tests round-trip write and load of instance state
+func TestJetStreamManager_WriteAndLoadState(t *testing.T) {
+	ns, natsURL := startTestNATSServerWithJetStream(t)
+	defer ns.Shutdown()
+
+	// Connect to NATS
+	nc, err := nats.Connect(natsURL)
+	require.NoError(t, err, "Failed to connect to NATS")
+	defer nc.Close()
+
+	// Create JetStreamManager
+	jsm, err := NewJetStreamManager(nc)
+	require.NoError(t, err, "Failed to create JetStreamManager")
+
+	// Initialize the KV bucket
+	err = jsm.InitKVBucket()
+	require.NoError(t, err, "Failed to init KV bucket")
+
+	// Create test instances
+	testNodeID := "test-node-1"
+	testInstances := &vm.Instances{
+		VMS: map[string]*vm.VM{
+			"i-test-001": {
+				ID:           "i-test-001",
+				Status:       "running",
+				InstanceType: "t3.micro",
+			},
+			"i-test-002": {
+				ID:           "i-test-002",
+				Status:       "stopped",
+				InstanceType: "t3.small",
+			},
+		},
+	}
+
+	// Write state
+	err = jsm.WriteState(testNodeID, testInstances)
+	require.NoError(t, err, "Failed to write state")
+
+	// Load state
+	loadedInstances, err := jsm.LoadState(testNodeID)
+	require.NoError(t, err, "Failed to load state")
+	require.NotNil(t, loadedInstances, "Loaded instances should not be nil")
+
+	// Verify the loaded state matches
+	assert.Len(t, loadedInstances.VMS, 2, "Should have 2 instances")
+	assert.NotNil(t, loadedInstances.VMS["i-test-001"], "Should have i-test-001")
+	assert.NotNil(t, loadedInstances.VMS["i-test-002"], "Should have i-test-002")
+	assert.Equal(t, "running", loadedInstances.VMS["i-test-001"].Status)
+	assert.Equal(t, "stopped", loadedInstances.VMS["i-test-002"].Status)
+	assert.Equal(t, "t3.micro", loadedInstances.VMS["i-test-001"].InstanceType)
+	assert.Equal(t, "t3.small", loadedInstances.VMS["i-test-002"].InstanceType)
+}
+
+// TestJetStreamManager_LoadState_KeyNotFound tests that LoadState returns empty state when key doesn't exist
+func TestJetStreamManager_LoadState_KeyNotFound(t *testing.T) {
+	ns, natsURL := startTestNATSServerWithJetStream(t)
+	defer ns.Shutdown()
+
+	nc, err := nats.Connect(natsURL)
+	require.NoError(t, err)
+	defer nc.Close()
+
+	jsm, err := NewJetStreamManager(nc)
+	require.NoError(t, err)
+
+	err = jsm.InitKVBucket()
+	require.NoError(t, err)
+
+	// Load state for a non-existent node
+	instances, err := jsm.LoadState("non-existent-node")
+	require.NoError(t, err, "Should not error when key not found")
+	require.NotNil(t, instances, "Should return non-nil instances")
+	assert.Empty(t, instances.VMS, "Should return empty VMS map")
+}
+
+// TestJetStreamManager_BucketCreation tests that InitKVBucket creates the bucket when it doesn't exist
+func TestJetStreamManager_BucketCreation(t *testing.T) {
+	ns, natsURL := startTestNATSServerWithJetStream(t)
+	defer ns.Shutdown()
+
+	nc, err := nats.Connect(natsURL)
+	require.NoError(t, err)
+	defer nc.Close()
+
+	jsm, err := NewJetStreamManager(nc)
+	require.NoError(t, err)
+
+	// InitKVBucket should create the bucket
+	err = jsm.InitKVBucket()
+	require.NoError(t, err, "Should create bucket without error")
+
+	// Verify the bucket exists by checking jsm.kv is set
+	assert.NotNil(t, jsm.kv, "KV bucket should be initialized")
+}
+
+// TestJetStreamManager_BucketReconnection tests that InitKVBucket connects to existing bucket
+func TestJetStreamManager_BucketReconnection(t *testing.T) {
+	ns, natsURL := startTestNATSServerWithJetStream(t)
+	defer ns.Shutdown()
+
+	// First connection - create the bucket
+	nc1, err := nats.Connect(natsURL)
+	require.NoError(t, err)
+
+	jsm1, err := NewJetStreamManager(nc1)
+	require.NoError(t, err)
+
+	err = jsm1.InitKVBucket()
+	require.NoError(t, err, "First InitKVBucket should succeed")
+
+	// Write some test data
+	testInstances := &vm.Instances{
+		VMS: map[string]*vm.VM{
+			"i-persist": {
+				ID:     "i-persist",
+				Status: "running",
+			},
+		},
+	}
+	err = jsm1.WriteState("persist-node", testInstances)
+	require.NoError(t, err)
+
+	nc1.Close()
+
+	// Second connection - should connect to existing bucket
+	nc2, err := nats.Connect(natsURL)
+	require.NoError(t, err)
+	defer nc2.Close()
+
+	jsm2, err := NewJetStreamManager(nc2)
+	require.NoError(t, err)
+
+	err = jsm2.InitKVBucket()
+	require.NoError(t, err, "Second InitKVBucket should succeed (reconnect)")
+
+	// Verify data persisted
+	loadedInstances, err := jsm2.LoadState("persist-node")
+	require.NoError(t, err)
+	assert.NotEmpty(t, loadedInstances.VMS, "Should have persisted instances")
+	assert.NotNil(t, loadedInstances.VMS["i-persist"], "Should have i-persist")
+	assert.Equal(t, "running", loadedInstances.VMS["i-persist"].Status)
+}
+
+// TestJetStreamManager_DeleteState tests deleting state from the KV store
+func TestJetStreamManager_DeleteState(t *testing.T) {
+	ns, natsURL := startTestNATSServerWithJetStream(t)
+	defer ns.Shutdown()
+
+	nc, err := nats.Connect(natsURL)
+	require.NoError(t, err)
+	defer nc.Close()
+
+	jsm, err := NewJetStreamManager(nc)
+	require.NoError(t, err)
+
+	err = jsm.InitKVBucket()
+	require.NoError(t, err)
+
+	// Write state
+	testNodeID := "delete-test-node"
+	testInstances := &vm.Instances{
+		VMS: map[string]*vm.VM{
+			"i-delete-me": {
+				ID:     "i-delete-me",
+				Status: "running",
+			},
+		},
+	}
+	err = jsm.WriteState(testNodeID, testInstances)
+	require.NoError(t, err)
+
+	// Verify state exists
+	loadedInstances, err := jsm.LoadState(testNodeID)
+	require.NoError(t, err)
+	assert.NotEmpty(t, loadedInstances.VMS)
+
+	// Delete state
+	err = jsm.DeleteState(testNodeID)
+	require.NoError(t, err, "Should delete state without error")
+
+	// Verify state is gone (should return empty state)
+	loadedInstances, err = jsm.LoadState(testNodeID)
+	require.NoError(t, err)
+	assert.Empty(t, loadedInstances.VMS, "Should return empty state after deletion")
+}
+
+// TestJetStreamManager_DeleteState_NonExistent tests deleting state that doesn't exist
+func TestJetStreamManager_DeleteState_NonExistent(t *testing.T) {
+	ns, natsURL := startTestNATSServerWithJetStream(t)
+	defer ns.Shutdown()
+
+	nc, err := nats.Connect(natsURL)
+	require.NoError(t, err)
+	defer nc.Close()
+
+	jsm, err := NewJetStreamManager(nc)
+	require.NoError(t, err)
+
+	err = jsm.InitKVBucket()
+	require.NoError(t, err)
+
+	// Delete non-existent state should not error
+	err = jsm.DeleteState("non-existent-node")
+	require.NoError(t, err, "Deleting non-existent state should not error")
+}
+
+// TestJetStreamManager_WriteState_UpdateExisting tests that writing state updates existing entry
+func TestJetStreamManager_WriteState_UpdateExisting(t *testing.T) {
+	ns, natsURL := startTestNATSServerWithJetStream(t)
+	defer ns.Shutdown()
+
+	nc, err := nats.Connect(natsURL)
+	require.NoError(t, err)
+	defer nc.Close()
+
+	jsm, err := NewJetStreamManager(nc)
+	require.NoError(t, err)
+
+	err = jsm.InitKVBucket()
+	require.NoError(t, err)
+
+	testNodeID := "update-test-node"
+
+	// Write initial state
+	initialInstances := &vm.Instances{
+		VMS: map[string]*vm.VM{
+			"i-initial": {
+				ID:     "i-initial",
+				Status: "running",
+			},
+		},
+	}
+	err = jsm.WriteState(testNodeID, initialInstances)
+	require.NoError(t, err)
+
+	// Update state with different instances
+	updatedInstances := &vm.Instances{
+		VMS: map[string]*vm.VM{
+			"i-initial": {
+				ID:     "i-initial",
+				Status: "stopped", // Changed status
+			},
+			"i-new": { // Added new instance
+				ID:     "i-new",
+				Status: "running",
+			},
+		},
+	}
+	err = jsm.WriteState(testNodeID, updatedInstances)
+	require.NoError(t, err)
+
+	// Load and verify updated state
+	loadedInstances, err := jsm.LoadState(testNodeID)
+	require.NoError(t, err)
+	assert.Len(t, loadedInstances.VMS, 2, "Should have 2 instances")
+	assert.Equal(t, "stopped", loadedInstances.VMS["i-initial"].Status, "Status should be updated")
+	assert.NotNil(t, loadedInstances.VMS["i-new"], "Should have new instance")
+}
+
+// TestJetStreamManager_MultipleNodes tests storing state for multiple nodes
+func TestJetStreamManager_MultipleNodes(t *testing.T) {
+	ns, natsURL := startTestNATSServerWithJetStream(t)
+	defer ns.Shutdown()
+
+	nc, err := nats.Connect(natsURL)
+	require.NoError(t, err)
+	defer nc.Close()
+
+	jsm, err := NewJetStreamManager(nc)
+	require.NoError(t, err)
+
+	err = jsm.InitKVBucket()
+	require.NoError(t, err)
+
+	// Write state for node-1
+	node1Instances := &vm.Instances{
+		VMS: map[string]*vm.VM{
+			"i-node1-001": {ID: "i-node1-001", Status: "running"},
+		},
+	}
+	err = jsm.WriteState("node-1", node1Instances)
+	require.NoError(t, err)
+
+	// Write state for node-2
+	node2Instances := &vm.Instances{
+		VMS: map[string]*vm.VM{
+			"i-node2-001": {ID: "i-node2-001", Status: "stopped"},
+			"i-node2-002": {ID: "i-node2-002", Status: "running"},
+		},
+	}
+	err = jsm.WriteState("node-2", node2Instances)
+	require.NoError(t, err)
+
+	// Load and verify node-1 state
+	loadedNode1, err := jsm.LoadState("node-1")
+	require.NoError(t, err)
+	assert.Len(t, loadedNode1.VMS, 1)
+	assert.NotNil(t, loadedNode1.VMS["i-node1-001"])
+
+	// Load and verify node-2 state
+	loadedNode2, err := jsm.LoadState("node-2")
+	require.NoError(t, err)
+	assert.Len(t, loadedNode2.VMS, 2)
+	assert.NotNil(t, loadedNode2.VMS["i-node2-001"])
+	assert.NotNil(t, loadedNode2.VMS["i-node2-002"])
+
+	// Verify node isolation - node-1 doesn't have node-2's instances
+	_, exists := loadedNode1.VMS["i-node2-001"]
+	assert.False(t, exists, "Node-1 should not have node-2's instances")
+}
+
+// TestJetStreamManager_KVNotInitialized tests error handling when KV is not initialized
+func TestJetStreamManager_KVNotInitialized(t *testing.T) {
+	ns, natsURL := startTestNATSServerWithJetStream(t)
+	defer ns.Shutdown()
+
+	nc, err := nats.Connect(natsURL)
+	require.NoError(t, err)
+	defer nc.Close()
+
+	// Create JetStreamManager but don't call InitKVBucket
+	jsm, err := NewJetStreamManager(nc)
+	require.NoError(t, err)
+
+	testInstances := &vm.Instances{VMS: make(map[string]*vm.VM)}
+	err = jsm.WriteState("test-node", testInstances)
+	assert.Error(t, err, "WriteState should error when KV not initialized")
+
+	_, err = jsm.LoadState("test-node")
+	assert.Error(t, err, "LoadState should error when KV not initialized")
+
+	err = jsm.DeleteState("test-node")
+	assert.Error(t, err, "DeleteState should error when KV not initialized")
+}

hive/services/nats/nats.go

@@ -79,6 +79,7 @@ func launchService(config *Config) (err error) {
 			Debug:      config.Debug,
 			LogFile:    config.LogFile,
 			JetStream:  config.JetStream,
+			StoreDir:   config.DataDir,
 		}
 
 		// Set defaults if not provided