diff --git a/pkg/chainstate/chainstate.go b/pkg/chainstate/chainstate.go
index 7926bb9..05231e7 100644
--- a/pkg/chainstate/chainstate.go
+++ b/pkg/chainstate/chainstate.go
@@ -36,6 +36,7 @@ func Load(path string) (*State, error) {
 		}
 		return nil, err
 	}
+
 	var state State
 	if err := json.Unmarshal(data, &state); err == nil {
 		if state.Projects == nil {
@@ -44,39 +45,10 @@ func Load(path string) (*State, error) {
 		return &state, nil
 	}
 
-	// Fallback to old format migration
-	var old map[string]int
-	if err := json.Unmarshal(data, &old); err == nil {
-		return migrateOldState(old), nil
-	}
-
-	// Unmarshal failed completely, return empty state
+	// Unmarshal failed, return empty state
 	return &State{Projects: make(map[string]*ProjectState)}, nil
 }
 
-func migrateOldState(old map[string]int) *State {
-	s := &State{Projects: make(map[string]*ProjectState)}
-	for key, value := range old {
-		parts := strings.SplitN(key, "|", 3)
-		if len(parts) == 2 {
-			// root|name
-			root := parts[0]
-			name := parts[1]
-			p := getOrCreateProject(s, root)
-			p.Chains[name] = value
-		} else if len(parts) == 3 && parts[1] == "@active" {
-			// root|@active|name
-			root := parts[0]
-			name := parts[2]
-			if value > 0 {
-				p := getOrCreateProject(s, root)
-				p.ActiveChain = name
-			}
-		}
-	}
-	return s
-}
-
 func getOrCreateProject(state *State, root string) *ProjectState {
 	cleanRoot := filepath.Clean(root)
 	p, ok := state.Projects[cleanRoot]
diff --git a/pkg/chainstate/chainstate_test.go b/pkg/chainstate/chainstate_test.go
new file mode 100644
index 0000000..f946725
--- /dev/null
+++ b/pkg/chainstate/chainstate_test.go
@@ -0,0 +1,106 @@
+package chainstate
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+func TestLoadSaveComplexState(t *testing.T) {
+	tmpDir := t.TempDir()
+	path := filepath.Join(tmpDir, "state.json")
+
+	state := &State{Projects: make(map[string]*ProjectState)}
+
+	// Set up some initial state
+	SetActive("/my/project1", "chainA", state)
+	SetStep("/my/project1", "chainA", 2, state)
+	SetStep("/my/project1", "chainB", 5, state)
+
+	SetActive("/my/project2", "chainX", state)
+
+	err := Save(path, state)
+	if err != nil {
+		t.Fatalf("Failed to save state: %v", err)
+	}
+
+	// Load it back
+	loaded, err := Load(path)
+	if err != nil {
+		t.Fatalf("Failed to load state: %v", err)
+	}
+
+	if ActiveName("/my/project1", loaded) != "chainA" {
+		t.Errorf("Expected active chain 'chainA', got '%s'", ActiveName("/my/project1", loaded))
+	}
+	if GetStep("/my/project1", "chainA", loaded) != 2 {
+		t.Errorf("Expected step 2, got %d", GetStep("/my/project1", "chainA", loaded))
+	}
+	if GetStep("/my/project1", "chainB", loaded) != 5 {
+		t.Errorf("Expected step 5, got %d", GetStep("/my/project1", "chainB", loaded))
+	}
+	if ActiveName("/my/project2", loaded) != "chainX" {
+		t.Errorf("Expected active chain 'chainX', got '%s'", ActiveName("/my/project2", loaded))
+	}
+}
+
+func TestClear(t *testing.T) {
+	state := &State{Projects: make(map[string]*ProjectState)}
+	SetStep("/root", "chain1", 1, state)
+	SetActive("/root", "chain1", state)
+	SetStep("/root", "chain2", 2, state)
+
+	// Clear specific chain that is active
+	Clear("/root", "chain1", state)
+
+	if GetStep("/root", "chain1", state) != 0 {
+		t.Errorf("Expected chain1 step to be 0")
+	}
+	if ActiveName("/root", state) != "" {
+		t.Errorf("Expected active name to be cleared")
+	}
+	if GetStep("/root", "chain2", state) != 2 {
+		t.Errorf("Expected chain2 step to remain 2")
+	}
+
+	// Clear whole project
+	Clear("/root", "", state)
+	if _, exists := state.Projects["/root"]; exists {
+		t.Errorf("Expected project /root to be completely removed")
+	}
+}
+
+func TestLoadEdgeCases(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	// 1. Missing file should return empty state
+	state, err := Load(filepath.Join(tmpDir, "missing.json"))
+	if err != nil {
+		t.Errorf("Missing file should not error, got: %v", err)
+	}
+	if state == nil || state.Projects == nil {
+		t.Errorf("Missing file should return initialized empty state")
+	}
+
+	// 2. Corrupt file should return empty state without crashing
+	path := filepath.Join(tmpDir, "corrupt.json")
+	os.WriteFile(path, []byte("{ invalid json"), 0o644)
+	state, err = Load(path)
+	if err != nil {
+		t.Errorf("Corrupt file should not error out completely, should recover empty state. Got error: %v", err)
+	}
+	if state == nil || state.Projects == nil {
+		t.Errorf("Corrupt file should return initialized empty state")
+	}
+
+	// 3. Empty file
+	pathEmpty := filepath.Join(tmpDir, "empty.json")
+	os.WriteFile(pathEmpty, []byte(""), 0o644)
+	state, err = Load(pathEmpty)
+	if err != nil {
+		t.Errorf("Empty file should not error, got: %v", err)
+	}
+	if state == nil || state.Projects == nil {
+		t.Errorf("Empty file should return initialized empty state")
+	}
+}