fix(checkpoint): address rollback undo review

Author: Yumiue

internal/checkpoint/per_edit_snapshot_test.go

@@ -1920,3 +1920,137 @@ func TestRestoreBaseline_ErrorsWhenPathMissingFromBaseline(t *testing.T) {
 		t.Fatalf("RestoreBaseline missing path error = %v", err)
 	}
 }
+
+func TestFinalizeExactForCheckpointPaths_CapturesSelectedCurrentPaths(t *testing.T) {
+	store, workdir := newTestStore(t)
+	targetA := writeWorkdirFile(t, workdir, "a.txt", "a before\n")
+	targetB := writeWorkdirFile(t, workdir, "b.txt", "b before\n")
+	if _, err := store.CapturePreWrite(targetA); err != nil {
+		t.Fatalf("CapturePreWrite(a): %v", err)
+	}
+	if _, err := store.CapturePreWrite(targetB); err != nil {
+		t.Fatalf("CapturePreWrite(b): %v", err)
+	}
+	if err := os.WriteFile(targetA, []byte("a source\n"), 0o644); err != nil {
+		t.Fatalf("write a source: %v", err)
+	}
+	if err := os.WriteFile(targetB, []byte("b source\n"), 0o644); err != nil {
+		t.Fatalf("write b source: %v", err)
+	}
+	if _, err := store.FinalizeWithExactState("cp-source"); err != nil {
+		t.Fatalf("FinalizeWithExactState: %v", err)
+	}
+	store.Reset()
+
+	if err := os.WriteFile(targetA, []byte("a guard\n"), 0o644); err != nil {
+		t.Fatalf("write a guard: %v", err)
+	}
+	if err := os.WriteFile(targetB, []byte("b guard\n"), 0o644); err != nil {
+		t.Fatalf("write b guard: %v", err)
+	}
+	written, err := store.FinalizeExactForCheckpointPaths("cp-guard", "cp-source", []string{"a.txt", "./b.txt", "a.txt"})
+	if err != nil {
+		t.Fatalf("FinalizeExactForCheckpointPaths: %v", err)
+	}
+	if !written {
+		t.Fatal("FinalizeExactForCheckpointPaths written = false, want true")
+	}
+
+	if err := os.WriteFile(targetA, []byte("a drift\n"), 0o644); err != nil {
+		t.Fatalf("write a drift: %v", err)
+	}
+	if err := os.WriteFile(targetB, []byte("b drift\n"), 0o644); err != nil {
+		t.Fatalf("write b drift: %v", err)
+	}
+	if err := store.RestoreExact(context.Background(), "cp-guard"); err != nil {
+		t.Fatalf("RestoreExact(cp-guard): %v", err)
+	}
+	if got := mustReadFile(t, targetA); got != "a guard\n" {
+		t.Fatalf("targetA = %q, want guard", got)
+	}
+	if got := mustReadFile(t, targetB); got != "b guard\n" {
+		t.Fatalf("targetB = %q, want guard", got)
+	}
+}
+
+func TestFinalizeExactForCheckpointPaths_CapturesDeletedAndCreatedCurrentState(t *testing.T) {
+	store, workdir := newTestStore(t)
+	deleted := writeWorkdirFile(t, workdir, "deleted.txt", "before delete\n")
+	created := filepath.Join(workdir, "created.txt")
+	if _, err := store.CapturePreWrite(deleted); err != nil {
+		t.Fatalf("CapturePreWrite(deleted): %v", err)
+	}
+	if _, err := store.CapturePreWrite(created); err != nil {
+		t.Fatalf("CapturePreWrite(created): %v", err)
+	}
+	if err := os.Remove(deleted); err != nil {
+		t.Fatalf("remove deleted current: %v", err)
+	}
+	if err := os.WriteFile(created, []byte("created current\n"), 0o644); err != nil {
+		t.Fatalf("write created current: %v", err)
+	}
+	if _, err := store.FinalizeWithExactState("cp-source"); err != nil {
+		t.Fatalf("FinalizeWithExactState: %v", err)
+	}
+	store.Reset()
+
+	written, err := store.FinalizeExactForCheckpointPaths("cp-guard", "cp-source", []string{"deleted.txt", "created.txt"})
+	if err != nil {
+		t.Fatalf("FinalizeExactForCheckpointPaths: %v", err)
+	}
+	if !written {
+		t.Fatal("FinalizeExactForCheckpointPaths written = false, want true")
+	}
+
+	if err := os.WriteFile(deleted, []byte("deleted drift\n"), 0o644); err != nil {
+		t.Fatalf("write deleted drift: %v", err)
+	}
+	if err := os.Remove(created); err != nil {
+		t.Fatalf("remove created drift: %v", err)
+	}
+	if err := store.RestoreExact(context.Background(), "cp-guard"); err != nil {
+		t.Fatalf("RestoreExact(cp-guard): %v", err)
+	}
+	if _, err := os.Stat(deleted); !os.IsNotExist(err) {
+		t.Fatalf("deleted should be absent after exact restore, stat err = %v", err)
+	}
+	if got := mustReadFile(t, created); got != "created current\n" {
+		t.Fatalf("created = %q, want current", got)
+	}
+}
+
+func TestFinalizeExactForCheckpointPaths_ValidatesInputsAndPaths(t *testing.T) {
+	store, workdir := newTestStore(t)
+	target := writeWorkdirFile(t, workdir, "tracked.txt", "before\n")
+	if _, err := store.CapturePreWrite(target); err != nil {
+		t.Fatalf("CapturePreWrite: %v", err)
+	}
+	if err := os.WriteFile(target, []byte("after\n"), 0o644); err != nil {
+		t.Fatalf("write after: %v", err)
+	}
+	if _, err := store.FinalizeWithExactState("cp-source"); err != nil {
+		t.Fatalf("FinalizeWithExactState: %v", err)
+	}
+	store.Reset()
+
+	tests := []struct {
+		name       string
+		checkpoint string
+		source     string
+		paths      []string
+		want       string
+	}{
+		{name: "empty checkpoint", checkpoint: "", source: "cp-source", paths: []string{"tracked.txt"}, want: "empty checkpointID"},
+		{name: "empty source", checkpoint: "cp-guard", source: "", paths: []string{"tracked.txt"}, want: "source checkpoint id required"},
+		{name: "empty paths", checkpoint: "cp-guard", source: "cp-source", paths: nil, want: "exact snapshot paths required"},
+		{name: "missing path", checkpoint: "cp-guard", source: "cp-source", paths: []string{"missing.txt"}, want: "baseline version for path missing.txt not found"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			_, err := store.FinalizeExactForCheckpointPaths(tt.checkpoint, tt.source, tt.paths)
+			if err == nil || !strings.Contains(err.Error(), tt.want) {
+				t.Fatalf("FinalizeExactForCheckpointPaths error = %v, want %q", err, tt.want)
+			}
+		})
+	}
+}

internal/runtime/checkpoint_flow_test.go

@@ -481,18 +481,25 @@ func TestUndoRestoreCheckpoint_RestoresBaselineRollbackGuardPaths(t *testing.T)
 	fixture := newRuntimeCheckpointFixture(t)
 	targetA := filepath.Join(fixture.workdir, "baseline-a.txt")
 	targetB := filepath.Join(fixture.workdir, "baseline-b.txt")
+	targetC := filepath.Join(fixture.workdir, "baseline-c.txt")
 	if err := os.WriteFile(targetA, []byte("a before"), 0o644); err != nil {
 		t.Fatalf("WriteFile(targetA before) error = %v", err)
 	}
 	if err := os.WriteFile(targetB, []byte("b before"), 0o644); err != nil {
 		t.Fatalf("WriteFile(targetB before) error = %v", err)
 	}
+	if err := os.WriteFile(targetC, []byte("c before"), 0o644); err != nil {
+		t.Fatalf("WriteFile(targetC before) error = %v", err)
+	}
 	if _, err := fixture.perEditStore.CapturePreWrite(targetA); err != nil {
 		t.Fatalf("CapturePreWrite(targetA) error = %v", err)
 	}
 	if _, err := fixture.perEditStore.CapturePreWrite(targetB); err != nil {
 		t.Fatalf("CapturePreWrite(targetB) error = %v", err)
 	}
+	if _, err := fixture.perEditStore.CapturePreWrite(targetC); err != nil {
+		t.Fatalf("CapturePreWrite(targetC) error = %v", err)
+	}
 
 	state := newRunState("run-baseline-undo", fixture.session)
 	if err := fixture.service.createStartOfTurnCheckpoint(context.Background(), &state); err != nil {
@@ -515,20 +522,26 @@ func TestUndoRestoreCheckpoint_RestoresBaselineRollbackGuardPaths(t *testing.T)
 	if err := os.WriteFile(targetB, []byte("b after"), 0o644); err != nil {
 		t.Fatalf("WriteFile(targetB after) error = %v", err)
 	}
+	if err := os.WriteFile(targetC, []byte("c after"), 0o644); err != nil {
+		t.Fatalf("WriteFile(targetC after) error = %v", err)
+	}
 
 	if _, err := fixture.service.RestoreCheckpoint(context.Background(), GatewayRestoreInput{
 		SessionID:    fixture.session.ID,
 		CheckpointID: cpRecord.CheckpointID,
 		Mode:         "baseline",
-		Paths:        []string{"baseline-a.txt"},
+		Paths:        []string{"baseline-a.txt", "baseline-b.txt"},
 	}); err != nil {
 		t.Fatalf("RestoreCheckpoint(baseline) error = %v", err)
 	}
 	if got := string(mustReadRuntimeFile(t, targetA)); got != "a before" {
 		t.Fatalf("baseline restored targetA = %q, want a before", got)
 	}
-	if got := string(mustReadRuntimeFile(t, targetB)); got != "b after" {
-		t.Fatalf("baseline restored targetB = %q, want b after", got)
+	if got := string(mustReadRuntimeFile(t, targetB)); got != "b before" {
+		t.Fatalf("baseline restored targetB = %q, want b before", got)
+	}
+	if got := string(mustReadRuntimeFile(t, targetC)); got != "c after" {
+		t.Fatalf("baseline restored targetC = %q, want c after", got)
 	}
 
 	if _, err := fixture.service.UndoRestoreCheckpoint(context.Background(), fixture.session.ID); err != nil {
@@ -540,6 +553,9 @@ func TestUndoRestoreCheckpoint_RestoresBaselineRollbackGuardPaths(t *testing.T)
 	if got := string(mustReadRuntimeFile(t, targetB)); got != "b after" {
 		t.Fatalf("undo targetB = %q, want b after", got)
 	}
+	if got := string(mustReadRuntimeFile(t, targetC)); got != "c after" {
+		t.Fatalf("undo targetC = %q, want c after", got)
+	}
 }
 
 func TestRestoreCheckpointBaselineRejectsPathsThatNormalizeEmpty(t *testing.T) {
@@ -606,6 +622,81 @@ func TestRestoreCheckpointBaselineWrapsRestoreBaselineError(t *testing.T) {
 	}
 }
 
+func TestRestoreCheckpointBaselineMarksGuardBrokenWhenRestoreFails(t *testing.T) {
+	fixture := newRuntimeCheckpointFixture(t)
+	target := filepath.Join(fixture.workdir, "baseline.txt")
+	if err := os.WriteFile(target, []byte("before baseline"), 0o644); err != nil {
+		t.Fatalf("WriteFile(before baseline) error = %v", err)
+	}
+	if _, err := fixture.perEditStore.CapturePreWrite(target); err != nil {
+		t.Fatalf("CapturePreWrite() error = %v", err)
+	}
+
+	state := newRunState("run-baseline-failed-restore", fixture.session)
+	if err := fixture.service.createStartOfTurnCheckpoint(context.Background(), &state); err != nil {
+		t.Fatalf("createStartOfTurnCheckpoint() error = %v", err)
+	}
+	records, err := fixture.checkpointStore.ListCheckpoints(context.Background(), fixture.session.ID, checkpoint.ListCheckpointOpts{})
+	if err != nil {
+		t.Fatalf("ListCheckpoints() error = %v", err)
+	}
+	if len(records) != 1 {
+		t.Fatalf("records = %#v, want 1", records)
+	}
+	cpRecord := records[0]
+	if err := fixture.checkpointStore.UpdateCheckpointStatus(context.Background(), cpRecord.CheckpointID, agentsession.CheckpointStatusAvailable); err != nil {
+		t.Fatalf("UpdateCheckpointStatus() error = %v", err)
+	}
+	if err := os.Remove(target); err != nil {
+		t.Fatalf("Remove(target) error = %v", err)
+	}
+	if err := os.Mkdir(target, 0o755); err != nil {
+		t.Fatalf("Mkdir(target) error = %v", err)
+	}
+
+	_, _, err = fixture.service.restoreCheckpointBaseline(context.Background(), fixture.session.ID, cpRecord.CheckpointID, []string{"baseline.txt"})
+	if err == nil || !strings.Contains(err.Error(), "baseline restore code") {
+		t.Fatalf("restoreCheckpointBaseline() error = %v, want baseline restore code", err)
+	}
+
+	records, err = fixture.checkpointStore.ListCheckpoints(context.Background(), fixture.session.ID, checkpoint.ListCheckpointOpts{})
+	if err != nil {
+		t.Fatalf("ListCheckpoints(all) error = %v", err)
+	}
+	seenBrokenGuard := false
+	for _, record := range records {
+		if record.Reason != agentsession.CheckpointReasonGuard {
+			continue
+		}
+		if record.Status != agentsession.CheckpointStatusBroken {
+			t.Fatalf("guard status = %q, want broken", record.Status)
+		}
+		if record.CodeCheckpointRef != "" {
+			if perEditID := checkpoint.PerEditCheckpointIDFromRef(record.CodeCheckpointRef); perEditID != "" {
+				if err := fixture.perEditStore.RestoreExact(context.Background(), perEditID); err == nil {
+					t.Fatalf("failed restore guard %q still has restorable per-edit metadata", perEditID)
+				}
+			}
+		}
+		seenBrokenGuard = true
+	}
+	if !seenBrokenGuard {
+		t.Fatal("expected failed baseline restore to leave a broken guard record")
+	}
+
+	available, err := fixture.checkpointStore.ListCheckpoints(context.Background(), fixture.session.ID, checkpoint.ListCheckpointOpts{
+		RestorableOnly: true,
+	})
+	if err != nil {
+		t.Fatalf("ListCheckpoints(restorable) error = %v", err)
+	}
+	for _, record := range available {
+		if record.Reason == agentsession.CheckpointReasonGuard {
+			t.Fatalf("broken guard %q should not be returned as restorable", record.CheckpointID)
+		}
+	}
+}
+
 func TestUndoRestoreCheckpoint_RestoresGuardState(t *testing.T) {
 	fixture := newRuntimeCheckpointFixture(t)
 	target := filepath.Join(fixture.workdir, "undo.txt")

internal/runtime/checkpoint_restore.go

@@ -372,6 +372,10 @@ func (s *Service) restoreCheckpointBaseline(
 		return RestoreResult{}, agentsession.CheckpointRecord{}, fmt.Errorf("checkpoint: create baseline guard: %w", guardErr)
 	}
 	if err := s.perEditStore.RestoreBaseline(ctx, perEditID, relPaths); err != nil {
+		if guardWritten {
+			_ = s.perEditStore.DeleteCheckpoint(guardID)
+		}
+		_ = s.checkpointStore.UpdateCheckpointStatus(ctx, guardRecord.CheckpointID, agentsession.CheckpointStatusBroken)
 		return RestoreResult{}, agentsession.CheckpointRecord{}, fmt.Errorf("checkpoint: baseline restore code: %w", err)
 	}
 	return RestoreResult{CheckpointID: checkpointID, SessionID: sessionID}, guardRecord, nil