Add support for tracking and persisting snapshot IDs

Introduce `currentSnapshotID` to track the active snapshot and persist its state in a file. This change ensures the current snapshot ID is restored during initialization and maintained consistently across snapshot operations.

api.go

@@ -36,10 +36,12 @@ type AgateOptions struct {
 
 // Agate is the main entry point for the snapshot library.
 type Agate struct {
-	manager     SnapshotManager
+	manager           SnapshotManager
-	options     AgateOptions
+	options           AgateOptions
-	metadataDir string
+	metadataDir       string
-	blobsDir    string
+	blobsDir          string
+	currentSnapshotID string
+	currentIDFile     string
 }
 
 // New initializes a new instance of the Agate library with the given options.
@@ -90,16 +92,38 @@ func New(options AgateOptions) (*Agate, error) {
 		return nil, fmt.Errorf("failed to create snapshot manager: %w", err)
 	}
 
-	return &Agate{
+	// Create a file path for storing the current snapshot ID
-		manager:     manager,
+	currentIDFile := filepath.Join(options.WorkDir, "current_snapshot_id")
-		options:     options,
+
-		metadataDir: metadataDir,
+	agate := &Agate{
-		blobsDir:    blobsDir,
+		manager:       manager,
-	}, nil
+		options:       options,
+		metadataDir:   metadataDir,
+		blobsDir:      blobsDir,
+		currentIDFile: currentIDFile,
+	}
+
+	// Load the current snapshot ID if it exists
+	if _, err := os.Stat(currentIDFile); err == nil {
+		data, err := os.ReadFile(currentIDFile)
+		if err == nil && len(data) > 0 {
+			agate.currentSnapshotID = string(data)
+		}
+	}
+
+	// Call OpenFunc if provided to initialize resources in the active directory
+	if options.OpenFunc != nil {
+		if err := options.OpenFunc(options.BlobStore.GetActiveDir()); err != nil {
+			return nil, fmt.Errorf("failed to open resources during initialization: %w", err)
+		}
+	}
+
+	return agate, nil
 }
 
-// SaveSnapshot creates a new snapshot from the current state of the work directory.
+// SaveSnapshot creates a new snapshot from the current state of the active directory.
 // If parentID is provided, it will be set as the parent of the new snapshot.
+// If parentID is empty, it will use the ID of the snapshot currently loaded in the active directory.
 // Returns the ID of the created snapshot.
 func (a *Agate) SaveSnapshot(ctx context.Context, name string, parentID string) (string, error) {
 	// Call CloseFunc if provided
@@ -109,15 +133,26 @@ func (a *Agate) SaveSnapshot(ctx context.Context, name string, parentID string)
 		}
 	}
 
+	// If parentID is not provided, use the current snapshot ID
+	effectiveParentID := parentID
+
 	// Create the snapshot
-	snapshot, err := a.manager.CreateSnapshot(ctx, a.options.WorkDir, name, parentID)
+	snapshot, err := a.manager.CreateSnapshot(ctx, a.options.BlobStore.GetActiveDir(), name, effectiveParentID)
 	if err != nil {
 		return "", fmt.Errorf("failed to create snapshot: %w", err)
 	}
 
+	// Update the current snapshot ID to the newly created snapshot
+	a.currentSnapshotID = snapshot.ID
+
+	// Save the current snapshot ID to a file
+	if err := a.saveCurrentSnapshotID(); err != nil {
+		return "", fmt.Errorf("failed to save current snapshot ID: %w", err)
+	}
+
 	// Call OpenFunc if provided
 	if a.options.OpenFunc != nil {
-		if err := a.options.OpenFunc(a.options.WorkDir); err != nil {
+		if err := a.options.OpenFunc(a.options.BlobStore.GetActiveDir()); err != nil {
 			return "", fmt.Errorf("failed to open resources after snapshot: %w", err)
 		}
 	}
@@ -125,7 +160,7 @@ func (a *Agate) SaveSnapshot(ctx context.Context, name string, parentID string)
 	return snapshot.ID, nil
 }
 
-// RestoreSnapshot extracts a snapshot to the work directory.
+// RestoreSnapshot extracts a snapshot to the active directory.
 func (a *Agate) RestoreSnapshot(ctx context.Context, snapshotID string) error {
 	// Call CloseFunc if provided
 	if a.options.CloseFunc != nil {
@@ -135,13 +170,55 @@ func (a *Agate) RestoreSnapshot(ctx context.Context, snapshotID string) error {
 	}
 
 	// Extract the snapshot
-	if err := a.manager.ExtractSnapshot(ctx, snapshotID, a.options.WorkDir); err != nil {
+	if err := a.manager.ExtractSnapshot(ctx, snapshotID, a.options.BlobStore.GetActiveDir()); err != nil {
 		return fmt.Errorf("failed to extract snapshot: %w", err)
 	}
 
+	// Save the ID of the snapshot that was restored
+	a.currentSnapshotID = snapshotID
+
+	// Save the current snapshot ID to a file
+	if err := a.saveCurrentSnapshotID(); err != nil {
+		return fmt.Errorf("failed to save current snapshot ID: %w", err)
+	}
+
 	// Call OpenFunc if provided
 	if a.options.OpenFunc != nil {
-		if err := a.options.OpenFunc(a.options.WorkDir); err != nil {
+		if err := a.options.OpenFunc(a.options.BlobStore.GetActiveDir()); err != nil {
+			return fmt.Errorf("failed to open resources after restore: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// RestoreSnapshot extracts a snapshot to the directory.
+func (a *Agate) RestoreSnapshotToDir(ctx context.Context, snapshotID string, dir string) error {
+	// Call CloseFunc if provided
+	if a.options.CloseFunc != nil {
+		if err := a.options.CloseFunc(); err != nil {
+			return fmt.Errorf("failed to close resources before restore: %w", err)
+		}
+	}
+
+	// Extract the snapshot
+	if err := a.manager.ExtractSnapshot(ctx, snapshotID, dir); err != nil {
+		return fmt.Errorf("failed to extract snapshot: %w", err)
+	}
+
+	// If restoring to the active directory, save the snapshot ID
+	if dir == a.options.BlobStore.GetActiveDir() {
+		a.currentSnapshotID = snapshotID
+
+		// Save the current snapshot ID to a file
+		if err := a.saveCurrentSnapshotID(); err != nil {
+			return fmt.Errorf("failed to save current snapshot ID: %w", err)
+		}
+	}
+
+	// Call OpenFunc if provided
+	if a.options.OpenFunc != nil {
+		if err := a.options.OpenFunc(dir); err != nil {
 			return fmt.Errorf("failed to open resources after restore: %w", err)
 		}
 	}
@@ -164,6 +241,20 @@ func (a *Agate) DeleteSnapshot(ctx context.Context, snapshotID string) error {
 	return a.manager.DeleteSnapshot(ctx, snapshotID)
 }
 
+// saveCurrentSnapshotID saves the current snapshot ID to a file in the WorkDir
+func (a *Agate) saveCurrentSnapshotID() error {
+	if a.currentSnapshotID == "" {
+		// If there's no current snapshot ID, remove the file if it exists
+		if _, err := os.Stat(a.currentIDFile); err == nil {
+			return os.Remove(a.currentIDFile)
+		}
+		return nil
+	}
+
+	// Write the current snapshot ID to the file
+	return os.WriteFile(a.currentIDFile, []byte(a.currentSnapshotID), 0644)
+}
+
 // Close releases all resources used by the Agate instance.
 func (a *Agate) Close() error {
 	// Currently, we don't have a way to close the manager directly

go.mod

@@ -6,14 +6,14 @@ require (
 	github.com/google/uuid v1.6.0
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.3
 	github.com/mattn/go-sqlite3 v1.14.28
-	google.golang.org/genproto/googleapis/api v0.0.0-20250422160041-2d3770c4ea7f
+	google.golang.org/genproto/googleapis/api v0.0.0-20250505200425-f936aa4a68b2
 	google.golang.org/grpc v1.72.0
 	google.golang.org/protobuf v1.36.6
 )
 
 require (
-	golang.org/x/net v0.39.0 // indirect
+	golang.org/x/net v0.40.0 // indirect
-	golang.org/x/sys v0.32.0 // indirect
+	golang.org/x/sys v0.33.0 // indirect
-	golang.org/x/text v0.24.0 // indirect
+	golang.org/x/text v0.25.0 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20250422160041-2d3770c4ea7f // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20250505200425-f936aa4a68b2 // indirect
 )

go.sum

@@ -26,14 +26,24 @@ go.opentelemetry.io/otel/trace v1.34.0 h1:+ouXS2V8Rd4hp4580a8q23bg0azF2nI8cqLYnC
 go.opentelemetry.io/otel/trace v1.34.0/go.mod h1:Svm7lSjQD7kG7KJ/MUHPVXSDGz2OX4h0M2jHBhmSfRE=
 golang.org/x/net v0.39.0 h1:ZCu7HMWDxpXpaiKdhzIfaltL9Lp31x/3fCP11bc6/fY=
 golang.org/x/net v0.39.0/go.mod h1:X7NRbYVEA+ewNkCNyJ513WmMdQ3BineSwVtN2zD/d+E=
+golang.org/x/net v0.40.0 h1:79Xs7wF06Gbdcg4kdCCIQArK11Z1hr5POQ6+fIYHNuY=
+golang.org/x/net v0.40.0/go.mod h1:y0hY0exeL2Pku80/zKK7tpntoX23cqL3Oa6njdgRtds=
 golang.org/x/sys v0.32.0 h1:s77OFDvIQeibCmezSnk/q6iAfkdiQaJi4VzroCFrN20=
 golang.org/x/sys v0.32.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
+golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw=
+golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
 golang.org/x/text v0.24.0 h1:dd5Bzh4yt5KYA8f9CJHCP4FB4D51c2c6JvN37xJJkJ0=
 golang.org/x/text v0.24.0/go.mod h1:L8rBsPeo2pSS+xqN0d5u2ikmjtmoJbDBT1b7nHvFCdU=
+golang.org/x/text v0.25.0 h1:qVyWApTSYLk/drJRO5mDlNYskwQznZmkpV2c8q9zls4=
+golang.org/x/text v0.25.0/go.mod h1:WEdwpYrmk1qmdHvhkSTNPm3app7v4rsT8F2UD6+VHIA=
 google.golang.org/genproto/googleapis/api v0.0.0-20250422160041-2d3770c4ea7f h1:tjZsroqekhC63+WMqzmWyW5Twj/ZfR5HAlpd5YQ1Vs0=
 google.golang.org/genproto/googleapis/api v0.0.0-20250422160041-2d3770c4ea7f/go.mod h1:Cd8IzgPo5Akum2c9R6FsXNaZbH3Jpa2gpHlW89FqlyQ=
+google.golang.org/genproto/googleapis/api v0.0.0-20250505200425-f936aa4a68b2 h1:vPV0tzlsK6EzEDHNNH5sa7Hs9bd7iXR7B1tSiPepkV0=
+google.golang.org/genproto/googleapis/api v0.0.0-20250505200425-f936aa4a68b2/go.mod h1:pKLAc5OolXC3ViWGI62vvC0n10CpwAtRcTNCFwTKBEw=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20250422160041-2d3770c4ea7f h1:N/PrbTw4kdkqNRzVfWPrBekzLuarFREcbFOiOLkXon4=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20250422160041-2d3770c4ea7f/go.mod h1:qQ0YXyHHx3XkvlzUtpXDkS29lDSafHMZBAZDc03LQ3A=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20250505200425-f936aa4a68b2 h1:IqsN8hx+lWLqlN+Sc3DoMy/watjofWiU8sRFgQ8fhKM=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20250505200425-f936aa4a68b2/go.mod h1:qQ0YXyHHx3XkvlzUtpXDkS29lDSafHMZBAZDc03LQ3A=
 google.golang.org/grpc v1.72.0 h1:S7UkcVa60b5AAQTaO6ZKamFp1zMZSU0fGDK2WZLbBnM=
 google.golang.org/grpc v1.72.0/go.mod h1:wH5Aktxcg25y1I3w7H69nHfXdOG3UiadoBtjh3izSDM=
 google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY=