atelet: prune local pause snapshots

cmd/atelet/local_checkpoints.go

@@ -0,0 +1,54 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"context"
+	"log/slog"
+	"os"
+	"path/filepath"
+
+	"github.com/agent-substrate/substrate/internal/ateompath"
+)
+
+// pruneLocalCheckpoints removes the actor's local pause snapshots except keep;
+// empty keep removes all. Best-effort: failures are logged, never fatal.
+func pruneLocalCheckpoints(ctx context.Context, actorUID, keep string) {
+	pruneLocalCheckpointDir(ctx, ateompath.LocalCheckpointsDir(actorUID), keep)
+}
+
+func pruneLocalCheckpointDir(ctx context.Context, dir, keep string) {
+	entries, err := os.ReadDir(dir)
+	if err != nil {
+		if !os.IsNotExist(err) {
+			slog.WarnContext(ctx, "failed to list local checkpoints for pruning", slog.String("dir", dir), slog.Any("err", err))
+		}
+		return
+	}
+	for _, entry := range entries {
+		if keep != "" && entry.Name() == keep {
+			continue
+		}
+		path := filepath.Join(dir, entry.Name())
+		if err := os.RemoveAll(path); err != nil {
+			slog.WarnContext(ctx, "failed to prune local checkpoint", slog.String("path", path), slog.Any("err", err))
+			continue
+		}
+		slog.InfoContext(ctx, "pruned local checkpoint", slog.String("path", path))
+	}
+	if keep == "" {
+		_ = os.Remove(dir)
+	}
+}

cmd/atelet/local_checkpoints_test.go

@@ -0,0 +1,78 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+func writeSnapshotDir(t *testing.T, dir, prefix string) {
+	t.Helper()
+	p := filepath.Join(dir, prefix)
+	if err := os.MkdirAll(p, 0o700); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.WriteFile(filepath.Join(p, "memory.img"), []byte("x"), 0o600); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func listDir(t *testing.T, dir string) []string {
+	t.Helper()
+	entries, err := os.ReadDir(dir)
+	if os.IsNotExist(err) {
+		return nil
+	}
+	if err != nil {
+		t.Fatal(err)
+	}
+	var names []string
+	for _, e := range entries {
+		names = append(names, e.Name())
+	}
+	return names
+}
+
+func TestPruneKeepsOnlyNamedSnapshot(t *testing.T) {
+	dir := t.TempDir()
+	writeSnapshotDir(t, dir, "pause-1")
+	writeSnapshotDir(t, dir, "pause-2")
+	writeSnapshotDir(t, dir, "pause-3")
+
+	pruneLocalCheckpointDir(context.Background(), dir, "pause-3")
+
+	if got := listDir(t, dir); len(got) != 1 || got[0] != "pause-3" {
+		t.Fatalf("dir = %v, want only pause-3", got)
+	}
+}
+
+func TestPruneEmptyKeepRemovesEverything(t *testing.T) {
+	dir := t.TempDir()
+	writeSnapshotDir(t, dir, "pause-1")
+	writeSnapshotDir(t, dir, "pause-2")
+
+	pruneLocalCheckpointDir(context.Background(), dir, "")
+
+	if _, err := os.Stat(dir); !os.IsNotExist(err) {
+		t.Fatalf("dir still exists (err=%v), want removed entirely", err)
+	}
+}
+
+func TestPruneMissingDirIsNoop(t *testing.T) {
+	pruneLocalCheckpointDir(context.Background(), filepath.Join(t.TempDir(), "absent"), "anything")
+}

cmd/atelet/main.go

@@ -399,10 +399,14 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe
 		if err := s.uploadExternalCheckpoint(ctx, req, checkpointDir, sandboxRec); err != nil {
 			return nil, ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonFaileSaveSnapshot, ateerrors.ActorCrashedMetadata(), fmt.Errorf("%w: while uploading external snapshot: %w", ateerrors.ReasonFaileSaveSnapshot, err))
 		}
+		// The durable snapshot supersedes local pause snapshots.
+		pruneLocalCheckpoints(ctx, actorUID, "")
 	case ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL:
 		if err := s.moveLocalCheckpoint(ctx, req, checkpointDir, sandboxRec); err != nil {
 			return nil, ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonFaileSaveSnapshot, ateerrors.ActorCrashedMetadata(), fmt.Errorf("%w: while moving to local snapshot: %w", ateerrors.ReasonFaileSaveSnapshot, err))
 		}
+		// Without pruning, every pause leaves a snapshot behind until the disk fills.
+		pruneLocalCheckpoints(ctx, actorUID, req.GetLocalConfig().GetSnapshotPrefix())
 	default:
 		return nil, fmt.Errorf("unexpected checkpoint type: %v", req.GetType())
 	}
@@ -1049,8 +1053,8 @@ func validateCheckpointRequest(req *ateletpb.CheckpointRequest) error {
 			return err
 		}
 	case ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL:
-		if req.GetLocalConfig().GetSnapshotPrefix() == "" {
-			return fmt.Errorf("snapshot prefix must be non-empty for type %s", req.GetType().String())
+		if err := resources.ValidateLocalSnapshotPrefix(req.GetLocalConfig().GetSnapshotPrefix()); err != nil {
+			return err
 		}
 	default:
 		return fmt.Errorf("invalid checkpoint type: %v", req.GetType())
@@ -1101,8 +1105,8 @@ func validateRestoreRequest(req *ateletpb.RestoreRequest) error {
 			return err
 		}
 	case ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL:
-		if req.GetLocalConfig().GetSnapshotPrefix() == "" {
-			return fmt.Errorf("snapshot prefix must be non-empty for type %s", req.GetType().String())
+		if err := resources.ValidateLocalSnapshotPrefix(req.GetLocalConfig().GetSnapshotPrefix()); err != nil {
+			return err
 		}
 	default:
 		return fmt.Errorf("invalid checkpoint type: %v", req.GetType())

cmd/atelet/main_test.go

@@ -277,6 +277,14 @@ func TestValidateCheckpointRequest(t *testing.T) {
 			r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL
 			r.Config = &ateletpb.CheckpointRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: ""}}
 		}), true},
+		{"nested local snapshot prefix", makeReq(func(r *ateletpb.CheckpointRequest) {
+			r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL
+			r.Config = &ateletpb.CheckpointRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: "pause/2"}}
+		}), true},
+		{"traversal local snapshot prefix", makeReq(func(r *ateletpb.CheckpointRequest) {
+			r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL
+			r.Config = &ateletpb.CheckpointRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: ".."}}
+		}), true},
 		{"unspecified snapshot type", makeReq(func(r *ateletpb.CheckpointRequest) { r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_UNSPECIFIED }), true},
 		{"unspecified snapshot scope", makeReq(func(r *ateletpb.CheckpointRequest) { r.Scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_UNSPECIFIED }), true},
 		{"invalid snapshot scope", makeReq(func(r *ateletpb.CheckpointRequest) { r.Scope = ateletpb.SnapshotScope(23) }), true},
@@ -323,6 +331,14 @@ func TestValidateRestoreRequest(t *testing.T) {
 			r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL
 			r.Config = &ateletpb.RestoreRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: ""}}
 		}), true},
+		{"nested local snapshot prefix", makeReq(func(r *ateletpb.RestoreRequest) {
+			r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL
+			r.Config = &ateletpb.RestoreRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: "pause/2"}}
+		}), true},
+		{"traversal local snapshot prefix", makeReq(func(r *ateletpb.RestoreRequest) {
+			r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL
+			r.Config = &ateletpb.RestoreRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: ".."}}
+		}), true},
 		{"unspecified snapshot type", makeReq(func(r *ateletpb.RestoreRequest) { r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_UNSPECIFIED }), true},
 		{"unspecified snapshot scope", makeReq(func(r *ateletpb.RestoreRequest) { r.Scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_UNSPECIFIED }), true},
 		{"invalid snapshot scope", makeReq(func(r *ateletpb.RestoreRequest) { r.Scope = ateletpb.SnapshotScope(23) }), true},

internal/resources/validate.go

@@ -173,6 +173,20 @@ func ValidateSnapshotURIPrefix(prefix string) error {
 	return nil
 }
 
+// ValidateLocalSnapshotPrefix ensures a local snapshot prefix is a single
+// path segment: it is joined onto the actor's checkpoint directory and
+// compared against that directory's entries when pruning, so a nested or
+// relative prefix would escape the directory or delete the wrong snapshots.
+func ValidateLocalSnapshotPrefix(prefix string) error {
+	if prefix == "" {
+		return fmt.Errorf("snapshot prefix must be non-empty")
+	}
+	if prefix == "." || prefix == ".." || strings.ContainsAny(prefix, `/\`) {
+		return fmt.Errorf("invalid snapshot prefix %q: must be a single path segment", prefix)
+	}
+	return nil
+}
+
 // ValidateWorker checks that the worker message is well-formed.
 func ValidateWorker(worker *ateapipb.Worker, fldPath *field.Path) field.ErrorList {
 	var errs field.ErrorList

internal/resources/validate_test.go

@@ -238,6 +238,31 @@ func TestValidateSnapshotURIPrefix(t *testing.T) {
 	}
 }
 
+func TestValidateLocalSnapshotPrefix(t *testing.T) {
+	tests := []struct {
+		name    string
+		prefix  string
+		wantErr bool
+	}{
+		{"valid", "pause", false},
+		{"valid with dash and digits", "pause-2", false},
+		{"empty", "", true},
+		{"dot", ".", true},
+		{"dotdot", "..", true},
+		{"nested", "pause/2", true},
+		{"absolute", "/pause", true},
+		{"traversal", "../other-actor", true},
+		{"backslash", `pause\2`, true},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if err := ValidateLocalSnapshotPrefix(tt.prefix); (err != nil) != tt.wantErr {
+				t.Errorf("ValidateLocalSnapshotPrefix(%q) err = %v, wantErr %v", tt.prefix, err, tt.wantErr)
+			}
+		})
+	}
+}
+
 func TestValidateWorker(t *testing.T) {
 	tests := []struct {
 		name    string