consolidate forget tasks and add RepoOrchestrator interface for testability

Merge ScheduledForgetTask into taskforget.go, refactor Forget/ForgetAll into
a single Forget method with generic options, extract RepoOrchestrator interface
for task testing, and add comprehensive task run tests.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

internal/orchestrator/repo/repo.go

@@ -213,24 +213,22 @@ func (r *RepoOrchestrator) ListSnapshotFiles(ctx context.Context, snapshotId str
 	return lsEnts, nil
 }
 
-func (r *RepoOrchestrator) Forget(ctx context.Context, plan *v1.Plan, tags []string) ([]*v1.ResticSnapshot, error) {
+func (r *RepoOrchestrator) Forget(ctx context.Context, policy *v1.RetentionPolicy, opts ...restic.GenericOption) ([]*v1.ResticSnapshot, error) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	ctx, flush := forwardResticLogs(ctx)
 	defer flush()
 
-	policy := plan.Retention
 	if policy == nil {
-		return nil, fmt.Errorf("plan %q has no retention policy", plan.Id)
+		return nil, fmt.Errorf("repo %q: forget called with nil retention policy", r.repoConfig.Id)
 	}
 
 	result, err := r.repo.Forget(
-		ctx, protoutil.RetentionPolicyFromProto(plan.Retention),
-		restic.WithFlags("--tag", strings.Join(tags, ",")),
-		restic.WithFlags("--group-by", ""),
+		ctx, protoutil.RetentionPolicyFromProto(policy),
+		opts...,
 	)
 	if err != nil {
-		return nil, fmt.Errorf("forget snapshots for repo %v: %w", r.repoConfig.Id, err)
+		return nil, fmt.Errorf("forget for repo %v: %w", r.repoConfig.Id, err)
 	}
 
 	var forgotten []*v1.ResticSnapshot
@@ -242,43 +240,7 @@ func (r *RepoOrchestrator) Forget(ctx context.Context, plan *v1.Plan, tags []str
 		forgotten = append(forgotten, snapshotProto)
 	}
 
-	r.logger(ctx).Debug("forget snapshots", zap.String("plan", plan.Id), zap.Int("count", len(forgotten)), zap.Any("policy", policy))
-
-	return forgotten, nil
-}
-
-// ForgetAll runs forget with a retention policy applied to all snapshots in the repo,
-// grouped by tag (so each plan+instance combination is treated independently).
-func (r *RepoOrchestrator) ForgetAll(ctx context.Context, policy *v1.RetentionPolicy) ([]*v1.ResticSnapshot, error) {
-	r.mu.Lock()
-	defer r.mu.Unlock()
-	ctx, flush := forwardResticLogs(ctx)
-	defer flush()
-
-	if policy == nil {
-		return nil, fmt.Errorf("repo %q has no retention policy for scheduled forget", r.repoConfig.Id)
-	}
-
-	results, err := r.repo.ForgetAll(
-		ctx, protoutil.RetentionPolicyFromProto(policy),
-		restic.WithFlags("--group-by", "tag"),
-	)
-	if err != nil {
-		return nil, fmt.Errorf("scheduled forget for repo %v: %w", r.repoConfig.Id, err)
-	}
-
-	var forgotten []*v1.ResticSnapshot
-	for _, result := range results {
-		for _, snapshot := range result.Remove {
-			snapshotProto := protoutil.SnapshotToProto(&snapshot)
-			if err := protoutil.ValidateSnapshot(snapshotProto); err != nil {
-				return nil, fmt.Errorf("snapshot validation failed: %w", err)
-			}
-			forgotten = append(forgotten, snapshotProto)
-		}
-	}
-
-	r.logger(ctx).Debug("scheduled forget snapshots", zap.Int("count", len(forgotten)), zap.Any("policy", policy))
+	r.logger(ctx).Debug("forget snapshots", zap.Int("count", len(forgotten)), zap.Any("policy", policy))
 
 	return forgotten, nil
 }

internal/orchestrator/taskrunnerimpl.go

@@ -11,7 +11,6 @@ import (
 	"github.com/garethgeorge/backrest/internal/hook"
 	"github.com/garethgeorge/backrest/internal/oplog"
 	"github.com/garethgeorge/backrest/internal/orchestrator/logging"
-	"github.com/garethgeorge/backrest/internal/orchestrator/repo"
 	"github.com/garethgeorge/backrest/internal/orchestrator/tasks"
 	"github.com/google/uuid"
 	"go.uber.org/zap"
@@ -78,10 +77,6 @@ func (t *taskRunnerImpl) DeleteOperation(id ...int64) error {
 	return t.orchestrator.OpLog.Delete(id...)
 }
 
-func (t *taskRunnerImpl) Orchestrator() *Orchestrator {
-	return t.orchestrator
-}
-
 func (t *taskRunnerImpl) QueryOperations(q oplog.Query, fn func(*v1.Operation) error) error {
 	return t.orchestrator.OpLog.Query(q, fn)
 }
@@ -151,7 +146,7 @@ func (t *taskRunnerImpl) GetPlan(planID string) (*v1.Plan, error) {
 	return t.orchestrator.GetPlan(planID)
 }
 
-func (t *taskRunnerImpl) GetRepoOrchestrator(repoID string) (*repo.RepoOrchestrator, error) {
+func (t *taskRunnerImpl) GetRepoOrchestrator(repoID string) (tasks.RepoOrchestrator, error) {
 	return t.orchestrator.GetRepoOrchestrator(repoID)
 }
 

internal/orchestrator/tasks/task.go

@@ -10,11 +10,27 @@ import (
 	v1 "github.com/garethgeorge/backrest/gen/go/v1"
 	"github.com/garethgeorge/backrest/internal/config"
 	"github.com/garethgeorge/backrest/internal/oplog"
-	"github.com/garethgeorge/backrest/internal/orchestrator/repo"
+	"github.com/garethgeorge/backrest/pkg/restic"
 	"go.uber.org/zap"
 	"google.golang.org/protobuf/proto"
 )
 
+// RepoOrchestrator is the interface for repo operations that tasks depend on.
+// The concrete implementation is in the repo package.
+type RepoOrchestrator interface {
+	UnlockIfAutoEnabled(ctx context.Context) error
+	Backup(ctx context.Context, plan *v1.Plan, dryRun bool, progressCallback func(event *restic.BackupProgressEntry)) (*restic.BackupProgressEntry, error)
+	Forget(ctx context.Context, policy *v1.RetentionPolicy, opts ...restic.GenericOption) ([]*v1.ResticSnapshot, error)
+	ForgetSnapshot(ctx context.Context, snapshotId string) error
+	Prune(ctx context.Context, output io.Writer) error
+	Check(ctx context.Context, output io.Writer) error
+	Stats(ctx context.Context) (*v1.RepoStats, error)
+	Restore(ctx context.Context, snapshotId string, snapshotPath string, target string, progressCallback func(event *v1.RestoreProgressEntry)) (*v1.RestoreProgressEntry, error)
+	Snapshots(ctx context.Context) ([]*restic.Snapshot, error)
+	AddTags(ctx context.Context, snapshotIDs []string, tags []string) error
+	RunCommand(ctx context.Context, command string, writer io.Writer) error
+}
+
 var NeverScheduledTask = ScheduledTask{}
 
 const (
@@ -52,7 +68,7 @@ type TaskRunner interface {
 	// GetPlan returns the plan with the given ID.
 	GetPlan(planID string) (*v1.Plan, error)
 	// GetRepoOrchestrator returns the orchestrator for the repo with the given ID.
-	GetRepoOrchestrator(repoID string) (*repo.RepoOrchestrator, error)
+	GetRepoOrchestrator(repoID string) (RepoOrchestrator, error)
 	// ScheduleTask schedules a task to run at a specific time.
 	ScheduleTask(task Task, priority int) error
 	// Config returns the current config.
@@ -178,8 +194,24 @@ func curTimeMillis() int64 {
 }
 
 type testTaskRunner struct {
-	config *v1.Config // the config to use for the task runner.
+	config *v1.Config
 	oplog  *oplog.OpLog
+
+	// Configurable for Run() testing
+	orchestrator   RepoOrchestrator
+	hookCalls      []hookCall
+	scheduledTasks []scheduledTaskCall
+	onExecuteHooks func(ctx context.Context, events []v1.Hook_Condition, vars HookVars) error
+}
+
+type hookCall struct {
+	Events []v1.Hook_Condition
+	Vars   HookVars
+}
+
+type scheduledTaskCall struct {
+	Task     Task
+	Priority int
 }
 
 var _ TaskRunner = &testTaskRunner{}
@@ -224,7 +256,11 @@ func (t *testTaskRunner) DeleteOperation(id ...int64) error {
 }
 
 func (t *testTaskRunner) ExecuteHooks(ctx context.Context, events []v1.Hook_Condition, vars HookVars) error {
-	panic("not implemented")
+	t.hookCalls = append(t.hookCalls, hookCall{Events: events, Vars: vars})
+	if t.onExecuteHooks != nil {
+		return t.onExecuteHooks(ctx, events, vars)
+	}
+	return nil
 }
 
 func (t *testTaskRunner) QueryOperations(q oplog.Query, fn func(*v1.Operation) error) error {
@@ -250,12 +286,16 @@ func (t *testTaskRunner) GetPlan(planID string) (*v1.Plan, error) {
 	return cfg, nil
 }
 
-func (t *testTaskRunner) GetRepoOrchestrator(repoID string) (*repo.RepoOrchestrator, error) {
-	panic("not implemented")
+func (t *testTaskRunner) GetRepoOrchestrator(repoID string) (RepoOrchestrator, error) {
+	if t.orchestrator == nil {
+		return nil, errors.New("no repo orchestrator configured")
+	}
+	return t.orchestrator, nil
 }
 
 func (t *testTaskRunner) ScheduleTask(task Task, priority int) error {
-	panic("not implemented")
+	t.scheduledTasks = append(t.scheduledTasks, scheduledTaskCall{Task: task, Priority: priority})
+	return nil
 }
 
 func (t *testTaskRunner) Config() *v1.Config {
@@ -266,6 +306,12 @@ func (t *testTaskRunner) Logger(ctx context.Context) *zap.Logger {
 	return zap.L()
 }
 
+type nopWriteCloser struct {
+	io.Writer
+}
+
+func (nopWriteCloser) Close() error { return nil }
+
 func (t *testTaskRunner) LogrefWriter() (id string, w io.WriteCloser, err error) {
-	panic("not implemented")
+	return "test-logref", &nopWriteCloser{io.Discard}, nil
 }

internal/orchestrator/tasks/taskforget.go

@@ -2,23 +2,29 @@ package tasks
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"strings"
 	"time"
 
 	v1 "github.com/garethgeorge/backrest/gen/go/v1"
 	"github.com/garethgeorge/backrest/internal/oplog"
 	"github.com/garethgeorge/backrest/internal/orchestrator/repo"
+	"github.com/garethgeorge/backrest/internal/protoutil"
+	"github.com/garethgeorge/backrest/pkg/restic"
 	"github.com/hashicorp/go-multierror"
 	"go.uber.org/zap"
 )
 
-func NewOneoffForgetTask(repo *v1.Repo, planID string, flowID int64, at time.Time) Task {
+// NewOneoffForgetTask creates a per-plan forget task that runs once after a backup.
+// It applies the plan's retention policy scoped to snapshots tagged for that plan.
+func NewOneoffForgetTask(repoProto *v1.Repo, planID string, flowID int64, at time.Time) Task {
 	return &GenericOneoffTask{
 		OneoffTask: OneoffTask{
 			BaseTask: BaseTask{
 				TaskType:   "forget",
-				TaskName:   fmt.Sprintf("forget for plan %q in repo %q", planID, repo.Id),
-				TaskRepo:   repo,
+				TaskName:   fmt.Sprintf("forget for plan %q in repo %q", planID, repoProto.Id),
+				TaskRepo:   repoProto,
 				TaskPlanID: planID,
 			},
 			FlowID: flowID,
@@ -27,107 +33,273 @@ func NewOneoffForgetTask(repo *v1.Repo, planID string, flowID int64, at time.Tim
 				Op: &v1.Operation_OperationForget{},
 			},
 		},
-		Do: func(ctx context.Context, st ScheduledTask, taskRunner TaskRunner) error {
-			op := st.Op
-			forgetOp := op.GetOperationForget()
-			if forgetOp == nil {
+		Do: func(ctx context.Context, st ScheduledTask, runner TaskRunner) error {
+			if st.Op.GetOperationForget() == nil {
 				panic("forget task with non-forget operation")
 			}
 
-			return forgetHelper(ctx, st, taskRunner)
+			t := st.Task
+			l := runner.Logger(ctx)
+
+			plan, err := runner.GetPlan(t.PlanID())
+			if err != nil {
+				return fmt.Errorf("get plan %q: %w", t.PlanID(), err)
+			}
+
+			tags := []string{repo.TagForPlan(t.PlanID())}
+			if compat, err := UseLegacyCompatMode(l, runner, t.Repo().GetGuid(), t.PlanID()); err != nil {
+				return fmt.Errorf("check legacy compat mode: %w", err)
+			} else if !compat {
+				tags = append(tags, repo.TagForInstance(runner.Config().Instance))
+			} else {
+				l.Warn("forgetting snapshots without instance ID, using legacy behavior (e.g. --tags not including instance ID)")
+				l.Sugar().Warnf("to avoid this warning, tag all snapshots with the instance ID e.g. by running: \r\n"+
+					"restic tag --set '%s' --set '%s' --tag '%s'", repo.TagForPlan(t.PlanID()), repo.TagForInstance(runner.Config().Instance), repo.TagForPlan(t.PlanID()))
+			}
+
+			return forgetHelper(ctx, st, runner, plan.Retention,
+				restic.WithFlags("--tag", strings.Join(tags, ",")),
+				restic.WithFlags("--group-by", ""),
+			)
 		},
 	}
 }
 
-func forgetHelper(ctx context.Context, st ScheduledTask, taskRunner TaskRunner) error {
-	t := st.Task
-	l := taskRunner.Logger(ctx)
+// ScheduledForgetTask is a repo-level forget task that runs on a schedule.
+// It applies the repo's forget policy retention to all snapshots, grouped by tags.
+type ScheduledForgetTask struct {
+	BaseTask
+	force  bool
+	didRun bool
+}
 
-	// Helper to notify of errors
-	notifyError := func(err error) error {
-		return NotifyError(ctx, taskRunner, t.Name(), err, v1.Hook_CONDITION_FORGET_ERROR)
+func NewScheduledForgetTask(repoProto *v1.Repo, planID string, force bool) Task {
+	return &ScheduledForgetTask{
+		BaseTask: BaseTask{
+			TaskType:   "scheduled_forget",
+			TaskName:   fmt.Sprintf("scheduled forget for repo %q", repoProto.Id),
+			TaskRepo:   repoProto,
+			TaskPlanID: planID,
+		},
+		force: force,
+	}
+}
+
+func (t *ScheduledForgetTask) Next(now time.Time, runner TaskRunner) (ScheduledTask, error) {
+	if t.force {
+		if t.didRun {
+			return NeverScheduledTask, nil
+		}
+		t.didRun = true
+		return ScheduledTask{
+			Task:  t,
+			RunAt: now,
+			Op: &v1.Operation{
+				Op: &v1.Operation_OperationForget{},
+			},
+		}, nil
 	}
 
-	r, err := taskRunner.GetRepoOrchestrator(t.RepoID())
+	repoProto, err := runner.GetRepo(t.RepoID())
+	if err != nil {
+		return ScheduledTask{}, fmt.Errorf("get repo %v: %w", t.RepoID(), err)
+	}
+
+	if repoProto.GetForgetPolicy().GetSchedule() == nil {
+		return NeverScheduledTask, nil
+	}
+
+	var lastRan time.Time
+	var foundBackup bool
+	if err := runner.QueryOperations(oplog.Query{}.
+		SetInstanceID(runner.InstanceID()).
+		SetRepoGUID(repoProto.GetGuid()).
+		SetPlanID(PlanForSystemTasks).
+		SetReversed(true), func(op *v1.Operation) error {
+		if op.Status == v1.OperationStatus_STATUS_PENDING || op.Status == v1.OperationStatus_STATUS_SYSTEM_CANCELLED {
+			return nil
+		}
+		if _, ok := op.Op.(*v1.Operation_OperationForget); ok && op.UnixTimeEndMs != 0 {
+			lastRan = time.Unix(0, op.UnixTimeEndMs*int64(time.Millisecond))
+			return oplog.ErrStopIteration
+		}
+		if _, ok := op.Op.(*v1.Operation_OperationBackup); ok {
+			foundBackup = true
+		}
+		return nil
+	}); err != nil {
+		return NeverScheduledTask, fmt.Errorf("finding last scheduled forget run time: %w", err)
+	} else if !foundBackup {
+		lastRan = now
+	}
+
+	runAt, err := protoutil.ResolveSchedule(repoProto.GetForgetPolicy().GetSchedule(), lastRan, now)
+	if errors.Is(err, protoutil.ErrScheduleDisabled) {
+		return NeverScheduledTask, nil
+	} else if err != nil {
+		return NeverScheduledTask, fmt.Errorf("resolve schedule: %w", err)
+	}
+
+	return ScheduledTask{
+		Task:  t,
+		RunAt: runAt,
+		Op: &v1.Operation{
+			Op: &v1.Operation_OperationForget{},
+		},
+	}, nil
+}
+
+// shouldSkip returns true if there are no new successful backups since the last scheduled forget.
+func (t *ScheduledForgetTask) shouldSkip(runner TaskRunner, repoProto *v1.Repo) bool {
+	var lastForgetEndMs int64
+	var hasNewBackup bool
+
+	_ = runner.QueryOperations(oplog.Query{}.
+		SetInstanceID(runner.InstanceID()).
+		SetRepoGUID(repoProto.GetGuid()).
+		SetPlanID(PlanForSystemTasks).
+		SetReversed(true), func(op *v1.Operation) error {
+		if op.Status != v1.OperationStatus_STATUS_SUCCESS {
+			return nil
+		}
+		if _, ok := op.Op.(*v1.Operation_OperationForget); ok && op.UnixTimeEndMs != 0 {
+			lastForgetEndMs = op.UnixTimeEndMs
+			return oplog.ErrStopIteration
+		}
+		return nil
+	})
+
+	if lastForgetEndMs == 0 {
+		return false // no previous forget, don't skip
+	}
+
+	// Check if any backup completed after the last forget
+	_ = runner.QueryOperations(oplog.Query{}.
+		SetRepoGUID(repoProto.GetGuid()).
+		SetReversed(true), func(op *v1.Operation) error {
+		if op.UnixTimeEndMs < lastForgetEndMs {
+			return oplog.ErrStopIteration // older than last forget, stop looking
+		}
+		if op.Status == v1.OperationStatus_STATUS_SUCCESS {
+			if _, ok := op.Op.(*v1.Operation_OperationBackup); ok {
+				hasNewBackup = true
+				return oplog.ErrStopIteration
+			}
+		}
+		return nil
+	})
+
+	return !hasNewBackup
+}
+
+func (t *ScheduledForgetTask) Run(ctx context.Context, st ScheduledTask, runner TaskRunner) error {
+	op := st.Op
+
+	repoProto, err := runner.GetRepo(t.RepoID())
+	if err != nil {
+		return NotifyError(ctx, runner, t.Name(), fmt.Errorf("get repo %q: %w", t.RepoID(), err), v1.Hook_CONDITION_FORGET_ERROR)
+	}
+
+	// Skip if no new backups since last forget run.
+	// Mark as system-cancelled so it doesn't count as a successful run
+	// and the next schedule is computed from the last actual forget.
+	if t.shouldSkip(runner, repoProto) {
+		op.Op = &v1.Operation_OperationForget{
+			OperationForget: &v1.OperationForget{},
+		}
+		op.Status = v1.OperationStatus_STATUS_SYSTEM_CANCELLED
+		op.DisplayMessage = "Skipped: no new backups since last forget"
+		return nil
+	}
+
+	err = forgetHelper(ctx, st, runner, repoProto.GetForgetPolicy().GetRetention(),
+		restic.WithFlags("--group-by", "tags"),
+	)
+	if err != nil {
+		return err
+	}
+
+	// Schedule a stats task after successful forget
+	if e := runner.ScheduleTask(NewStatsTask(t.Repo(), PlanForSystemTasks, false), TaskPriorityStats); e != nil {
+		zap.L().Error("schedule stats task", zap.Error(e))
+	}
+
+	return nil
+}
+
+// forgetHelper contains the shared logic for running a forget operation.
+// It handles unlock, hooks, calling restic forget, and marking forgotten snapshots in the oplog.
+func forgetHelper(ctx context.Context, st ScheduledTask, runner TaskRunner, policy *v1.RetentionPolicy, opts ...restic.GenericOption) error {
+	t := st.Task
+
+	notifyError := func(err error) error {
+		return NotifyError(ctx, runner, t.Name(), err, v1.Hook_CONDITION_FORGET_ERROR)
+	}
+
+	r, err := runner.GetRepoOrchestrator(t.RepoID())
 	if err != nil {
 		return notifyError(fmt.Errorf("get repo %q: %w", t.RepoID(), err))
 	}
 
-	err = r.UnlockIfAutoEnabled(ctx)
-	if err != nil {
+	if err := r.UnlockIfAutoEnabled(ctx); err != nil {
 		return notifyError(fmt.Errorf("auto unlock repo %q: %w", t.RepoID(), err))
 	}
 
-	plan, err := taskRunner.GetPlan(t.PlanID())
-	if err != nil {
-		return notifyError(fmt.Errorf("get plan %q: %w", t.PlanID(), err))
-	}
-
-	// execute hooks
-	if err := taskRunner.ExecuteHooks(ctx, []v1.Hook_Condition{
+	if err := runner.ExecuteHooks(ctx, []v1.Hook_Condition{
 		v1.Hook_CONDITION_FORGET_START,
-	}, HookVars{Plan: plan}); err != nil {
+	}, HookVars{}); err != nil {
 		return notifyError(fmt.Errorf("forget start hook: %w", err))
 	}
 
-	tags := []string{repo.TagForPlan(t.PlanID())}
-	if compat, err := UseLegacyCompatMode(l, taskRunner, t.Repo().GetGuid(), t.PlanID()); err != nil {
-		return notifyError(fmt.Errorf("check legacy compat mode: %w", err))
-	} else if !compat {
-		tags = append(tags, repo.TagForInstance(taskRunner.Config().Instance))
-	} else {
-		l.Warn("forgetting snapshots without instance ID, using legacy behavior (e.g. --tags not including instance ID)")
-		l.Sugar().Warnf("to avoid this warning, tag all snapshots with the instance ID e.g. by running: \r\n"+
-			"restic tag --set '%s' --set '%s' --tag '%s'", repo.TagForPlan(t.PlanID()), repo.TagForInstance(taskRunner.Config().Instance), repo.TagForPlan(t.PlanID()))
-	}
-
-	// check if any other instance IDs exist in the repo (unassociated don't count)
-	forgot, err := r.Forget(ctx, plan, tags)
+	forgot, err := r.Forget(ctx, policy, opts...)
 
 	forgetOp := &v1.Operation_OperationForget{
-		OperationForget: &v1.OperationForget{},
+		OperationForget: &v1.OperationForget{
+			Forget: forgot,
+			Policy: policy,
+		},
 	}
 	st.Op.Op = forgetOp
 
-	forgetOp.OperationForget.Forget = append(forgetOp.OperationForget.Forget, forgot...)
-	forgetOp.OperationForget.Policy = plan.Retention
-
+	// Mark forgotten snapshots in the oplog
 	var ops []*v1.Operation
-	for _, forgot := range forgot {
-		if e := taskRunner.QueryOperations(oplog.Query{}.
+	for _, f := range forgot {
+		if e := runner.QueryOperations(oplog.Query{}.
 			SetRepoGUID(t.Repo().GetGuid()).
-			SetSnapshotID(forgot.Id), func(op *v1.Operation) error {
+			SetSnapshotID(f.Id), func(op *v1.Operation) error {
 			ops = append(ops, op)
 			return nil
 		}); e != nil {
-			err = multierror.Append(err, fmt.Errorf("cleanup snapshot %v: %w", forgot.Id, e))
+			err = multierror.Append(err, fmt.Errorf("lookup snapshot %v: %w", f.Id, e))
 		}
 	}
 
+	l := runner.Logger(ctx)
 	l.Sugar().Debugf("found %v snapshots were forgotten, marking this in oplog", len(ops))
 
 	for _, op := range ops {
 		if indexOp, ok := op.Op.(*v1.Operation_OperationIndexSnapshot); ok {
 			indexOp.OperationIndexSnapshot.Forgot = true
-			if e := taskRunner.UpdateOperation(op); err != nil {
+			if e := runner.UpdateOperation(op); e != nil {
 				err = multierror.Append(err, fmt.Errorf("mark index snapshot %v as forgotten: %w", op.Id, e))
-				continue
 			}
 		}
 	}
 
 	if err != nil {
 		return notifyError(fmt.Errorf("forget: %w", err))
-	} else if e := taskRunner.ExecuteHooks(ctx, []v1.Hook_Condition{
+	}
+
+	if e := runner.ExecuteHooks(ctx, []v1.Hook_Condition{
 		v1.Hook_CONDITION_FORGET_SUCCESS,
 	}, HookVars{}); e != nil {
 		return fmt.Errorf("forget end hook: %w", e)
 	}
 
-	return err
+	return nil
 }
 
-// useLegacyCompatMode checks if there are any snapshots that were created without a `created-by` tag still exist in the repo.
+// UseLegacyCompatMode checks if there are any snapshots that were created without a `created-by` tag still exist in the repo.
 // The property is overridden if mixed `created-by` tag values are found.
 func UseLegacyCompatMode(l *zap.Logger, taskRunner TaskRunner, repoGUID, planID string) (bool, error) {
 	instanceIDs := make(map[string]struct{})

internal/orchestrator/tasks/taskindexsnapshots.go

@@ -3,15 +3,12 @@ package tasks
 import (
 	"context"
 	"fmt"
-	"slices"
-	"strings"
 	"time"
 
 	v1 "github.com/garethgeorge/backrest/gen/go/v1"
 	"github.com/garethgeorge/backrest/internal/oplog"
 	"github.com/garethgeorge/backrest/internal/orchestrator/repo"
 	"github.com/garethgeorge/backrest/internal/protoutil"
-	"github.com/garethgeorge/backrest/pkg/restic"
 	"go.uber.org/zap"
 )
 
@@ -166,42 +163,3 @@ func instanceIDForSnapshot(snapshot *v1.ResticSnapshot) string {
 	return InstanceIDForUnassociatedOperations
 }
 
-// tryMigrate checks if the snapshots use the latest backrest tag set and migrates them if necessary.
-func tryMigrate(ctx context.Context, repo *repo.RepoOrchestrator, config *v1.Config, snapshots []*restic.Snapshot) (bool, error) {
-	if config.Instance == "" {
-		zap.S().Warnf("Instance ID not set. Skipping migration.")
-		return false, nil
-	}
-
-	planIDs := make(map[string]struct{})
-	for _, plan := range config.Plans {
-		planIDs[plan.Id] = struct{}{}
-	}
-
-	needsCreatedBy := []string{}
-	for _, snapshot := range snapshots {
-		// Check if snapshot is already tagged with `created-by:``
-		if idx := slices.IndexFunc(snapshot.Tags, func(tag string) bool {
-			return strings.HasPrefix(tag, "created-by:")
-		}); idx != -1 {
-			continue
-		}
-		// Check that snapshot is included in a plan for this instance. Backrest will not take ownership of snapshots belonging to it isn't aware of.
-		if _, ok := planIDs[planForSnapshot(protoutil.SnapshotToProto(snapshot))]; !ok {
-			continue
-		}
-		needsCreatedBy = append(needsCreatedBy, snapshot.Id)
-	}
-
-	if len(needsCreatedBy) == 0 {
-		return false, nil
-	}
-
-	zap.S().Warnf("Found %v snapshots without created-by tag but included in a plan for this instance. Taking ownership and adding created-by tag.", len(needsCreatedBy))
-
-	if err := repo.AddTags(ctx, needsCreatedBy, []string{fmt.Sprintf("created-by:%v", config.Instance)}); err != nil {
-		return false, fmt.Errorf("add created-by tag to snapshots: %w", err)
-	}
-
-	return true, nil
-}

internal/orchestrator/tasks/taskrun_test.go

@@ -0,0 +1,648 @@
+package tasks
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	v1 "github.com/garethgeorge/backrest/gen/go/v1"
+	"github.com/garethgeorge/backrest/internal/oplog"
+	"github.com/garethgeorge/backrest/internal/oplog/sqlitestore"
+	"github.com/garethgeorge/backrest/pkg/restic"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+const testSnapshotID = "1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef"
+
+func newTestConfig(repo *v1.Repo, plans ...*v1.Plan) *v1.Config {
+	return &v1.Config{
+		Instance: "test-instance",
+		Repos:    []*v1.Repo{repo},
+		Plans:    plans,
+	}
+}
+
+func setupTestRunner(t *testing.T, cfg *v1.Config, fake *fakeRepoOrchestrator) *testTaskRunner {
+	t.Helper()
+	opstore, err := sqlitestore.NewMemorySqliteStore(t)
+	require.NoError(t, err)
+	ol, err := oplog.NewOpLog(opstore)
+	require.NoError(t, err)
+	runner := newTestTaskRunner(t, cfg, ol)
+	runner.orchestrator = fake
+	return runner
+}
+
+func nextAndCreate(t *testing.T, task Task, runner *testTaskRunner) ScheduledTask {
+	t.Helper()
+	st, err := task.Next(time.Now(), runner)
+	require.NoError(t, err)
+	st.Task = task
+	if st.Op != nil {
+		// Populate fields the orchestrator normally sets before storing.
+		if st.Op.RepoId == "" && task.Repo() != nil {
+			st.Op.RepoId = task.Repo().Id
+		}
+		if st.Op.RepoGuid == "" && task.Repo() != nil {
+			st.Op.RepoGuid = task.Repo().Guid
+		}
+		if st.Op.PlanId == "" {
+			st.Op.PlanId = task.PlanID()
+		}
+		if st.Op.InstanceId == "" {
+			st.Op.InstanceId = runner.InstanceID()
+		}
+		if st.Op.FlowId == 0 {
+			st.Op.FlowId = 1
+		}
+		if st.Op.UnixTimeStartMs == 0 {
+			st.Op.UnixTimeStartMs = time.Now().UnixMilli()
+		}
+		require.NoError(t, runner.CreateOperation(st.Op))
+	}
+	return st
+}
+
+func hookContains(calls []hookCall, cond v1.Hook_Condition) bool {
+	for _, c := range calls {
+		for _, e := range c.Events {
+			if e == cond {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// --- PruneTask tests ---
+
+func TestPruneTaskRun(t *testing.T) {
+	tests := []struct {
+		name            string
+		fake            *fakeRepoOrchestrator
+		wantErr         bool
+		wantHooks       []v1.Hook_Condition
+		wantScheduled   int
+		scheduledType   string
+	}{
+		{
+			name:          "success",
+			fake:          &fakeRepoOrchestrator{},
+			wantHooks:     []v1.Hook_Condition{v1.Hook_CONDITION_PRUNE_START, v1.Hook_CONDITION_PRUNE_SUCCESS},
+			wantScheduled: 1,
+			scheduledType: "stats",
+		},
+		{
+			name:      "prune error",
+			fake:      &fakeRepoOrchestrator{pruneErr: fmt.Errorf("prune failed")},
+			wantErr:   true,
+			wantHooks: []v1.Hook_Condition{v1.Hook_CONDITION_PRUNE_START, v1.Hook_CONDITION_PRUNE_ERROR, v1.Hook_CONDITION_ANY_ERROR},
+		},
+		{
+			name:      "unlock error",
+			fake:      &fakeRepoOrchestrator{unlockErr: fmt.Errorf("unlock failed")},
+			wantErr:   true,
+			wantHooks: []v1.Hook_Condition{v1.Hook_CONDITION_PRUNE_ERROR, v1.Hook_CONDITION_ANY_ERROR},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			repo := &v1.Repo{Id: "repo1", Guid: "guid1"}
+			cfg := newTestConfig(repo)
+			runner := setupTestRunner(t, cfg, tc.fake)
+
+			task := NewPruneTask(repo, PlanForSystemTasks, true)
+			st := nextAndCreate(t, task, runner)
+
+			err := task.Run(context.Background(), st, runner)
+			if tc.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+
+			for _, cond := range tc.wantHooks {
+				assert.True(t, hookContains(runner.hookCalls, cond), "expected hook %v", cond)
+			}
+
+			assert.Len(t, runner.scheduledTasks, tc.wantScheduled)
+			if tc.scheduledType != "" && len(runner.scheduledTasks) > 0 {
+				assert.Equal(t, tc.scheduledType, runner.scheduledTasks[0].Task.Type())
+			}
+		})
+	}
+}
+
+// --- CheckTask tests ---
+
+func TestCheckTaskRun(t *testing.T) {
+	tests := []struct {
+		name      string
+		fake      *fakeRepoOrchestrator
+		wantErr   bool
+		wantHooks []v1.Hook_Condition
+	}{
+		{
+			name:      "success",
+			fake:      &fakeRepoOrchestrator{},
+			wantHooks: []v1.Hook_Condition{v1.Hook_CONDITION_CHECK_START, v1.Hook_CONDITION_CHECK_SUCCESS},
+		},
+		{
+			name:      "check error",
+			fake:      &fakeRepoOrchestrator{checkErr: fmt.Errorf("check failed")},
+			wantErr:   true,
+			wantHooks: []v1.Hook_Condition{v1.Hook_CONDITION_CHECK_START, v1.Hook_CONDITION_CHECK_ERROR, v1.Hook_CONDITION_ANY_ERROR},
+		},
+		{
+			name:      "unlock error",
+			fake:      &fakeRepoOrchestrator{unlockErr: fmt.Errorf("unlock failed")},
+			wantErr:   true,
+			wantHooks: []v1.Hook_Condition{v1.Hook_CONDITION_CHECK_ERROR, v1.Hook_CONDITION_ANY_ERROR},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			repo := &v1.Repo{Id: "repo1", Guid: "guid1"}
+			cfg := newTestConfig(repo)
+			runner := setupTestRunner(t, cfg, tc.fake)
+
+			task := NewCheckTask(repo, PlanForSystemTasks, true)
+			st := nextAndCreate(t, task, runner)
+
+			err := task.Run(context.Background(), st, runner)
+			if tc.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+
+			for _, cond := range tc.wantHooks {
+				assert.True(t, hookContains(runner.hookCalls, cond), "expected hook %v", cond)
+			}
+		})
+	}
+}
+
+// --- StatsTask tests ---
+
+func TestStatsTaskRun(t *testing.T) {
+	tests := []struct {
+		name    string
+		fake    *fakeRepoOrchestrator
+		wantErr bool
+	}{
+		{
+			name: "success",
+			fake: &fakeRepoOrchestrator{
+				statsResult: &v1.RepoStats{
+					TotalSize: 1000,
+				},
+			},
+		},
+		{
+			name:    "stats error",
+			fake:    &fakeRepoOrchestrator{statsErr: fmt.Errorf("stats failed")},
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			repo := &v1.Repo{Id: "repo1", Guid: "guid1"}
+			cfg := newTestConfig(repo)
+			runner := setupTestRunner(t, cfg, tc.fake)
+
+			task := NewStatsTask(repo, PlanForSystemTasks, true)
+			st := nextAndCreate(t, task, runner)
+
+			err := task.Run(context.Background(), st, runner)
+			if tc.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				statsOp := st.Op.GetOperationStats()
+				require.NotNil(t, statsOp)
+				assert.Equal(t, tc.fake.statsResult.TotalSize, statsOp.Stats.TotalSize)
+			}
+		})
+	}
+}
+
+// --- BackupTask tests ---
+
+func TestBackupTaskRun(t *testing.T) {
+	tests := []struct {
+		name          string
+		dryRun        bool
+		fake          *fakeRepoOrchestrator
+		repo          *v1.Repo
+		plan          *v1.Plan
+		wantErr       bool
+		wantHooks     []v1.Hook_Condition
+		wantNotHooks  []v1.Hook_Condition
+		wantScheduled []string // expected scheduled task types
+	}{
+		{
+			name: "successful backup with retention",
+			fake: &fakeRepoOrchestrator{
+				backupResult: &restic.BackupProgressEntry{
+					MessageType:         "summary",
+					SnapshotId:          testSnapshotID,
+					TotalBytesProcessed: 1000,
+				},
+			},
+			plan: &v1.Plan{
+				Id:   "plan1",
+				Repo: "repo1",
+				Retention: &v1.RetentionPolicy{
+					Policy: &v1.RetentionPolicy_PolicyKeepLastN{PolicyKeepLastN: 5},
+				},
+			},
+			wantHooks:     []v1.Hook_Condition{v1.Hook_CONDITION_SNAPSHOT_START, v1.Hook_CONDITION_SNAPSHOT_SUCCESS, v1.Hook_CONDITION_SNAPSHOT_END},
+			wantScheduled: []string{"forget", "index_snapshots"},
+		},
+		{
+			name: "successful backup no retention",
+			fake: &fakeRepoOrchestrator{
+				backupResult: &restic.BackupProgressEntry{
+					MessageType:         "summary",
+					SnapshotId:          testSnapshotID,
+					TotalBytesProcessed: 1000,
+				},
+			},
+			plan: &v1.Plan{
+				Id:   "plan1",
+				Repo: "repo1",
+			},
+			wantHooks:     []v1.Hook_Condition{v1.Hook_CONDITION_SNAPSHOT_START, v1.Hook_CONDITION_SNAPSHOT_SUCCESS, v1.Hook_CONDITION_SNAPSHOT_END},
+			wantScheduled: []string{"index_snapshots"},
+		},
+		{
+			name: "successful backup with repo-level scheduled forget skips per-plan forget",
+			repo: &v1.Repo{
+				Id: "repo1", Guid: "guid1",
+				ForgetPolicy: &v1.ForgetPolicy{
+					Schedule: &v1.Schedule{
+						Schedule: &v1.Schedule_MaxFrequencyDays{MaxFrequencyDays: 1},
+					},
+				},
+			},
+			fake: &fakeRepoOrchestrator{
+				backupResult: &restic.BackupProgressEntry{
+					MessageType: "summary",
+					SnapshotId:  testSnapshotID,
+				},
+			},
+			plan: &v1.Plan{
+				Id:   "plan1",
+				Repo: "repo1",
+				Retention: &v1.RetentionPolicy{
+					Policy: &v1.RetentionPolicy_PolicyKeepLastN{PolicyKeepLastN: 5},
+				},
+			},
+			wantScheduled: []string{"index_snapshots"}, // no forget
+		},
+		{
+			name:    "backup error",
+			fake:    &fakeRepoOrchestrator{backupErr: fmt.Errorf("backup failed")},
+			plan:    &v1.Plan{Id: "plan1", Repo: "repo1"},
+			wantErr: true,
+			wantHooks: []v1.Hook_Condition{
+				v1.Hook_CONDITION_SNAPSHOT_START,
+				v1.Hook_CONDITION_SNAPSHOT_ERROR,
+				v1.Hook_CONDITION_ANY_ERROR,
+				v1.Hook_CONDITION_SNAPSHOT_END,
+			},
+		},
+		{
+			name:    "unlock error",
+			fake:    &fakeRepoOrchestrator{unlockErr: fmt.Errorf("unlock failed")},
+			plan:    &v1.Plan{Id: "plan1", Repo: "repo1"},
+			wantErr: true,
+			wantHooks: []v1.Hook_Condition{
+				v1.Hook_CONDITION_SNAPSHOT_ERROR,
+				v1.Hook_CONDITION_ANY_ERROR,
+			},
+		},
+		{
+			name:   "dry run backup",
+			dryRun: true,
+			fake: &fakeRepoOrchestrator{
+				backupResult: &restic.BackupProgressEntry{
+					MessageType: "summary",
+					SnapshotId:  testSnapshotID,
+				},
+			},
+			plan:          &v1.Plan{Id: "plan1", Repo: "repo1"},
+			wantScheduled: nil,
+		},
+		{
+			name: "skip if unchanged",
+			fake: &fakeRepoOrchestrator{
+				backupResult: &restic.BackupProgressEntry{
+					MessageType: "summary",
+					SnapshotId:  "", // empty = no changes
+				},
+			},
+			plan:          &v1.Plan{Id: "plan1", Repo: "repo1"},
+			wantHooks:     []v1.Hook_Condition{v1.Hook_CONDITION_SNAPSHOT_START, v1.Hook_CONDITION_SNAPSHOT_SKIPPED, v1.Hook_CONDITION_SNAPSHOT_END},
+			wantScheduled: nil,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			repo := tc.repo
+			if repo == nil {
+				repo = &v1.Repo{Id: "repo1", Guid: "guid1"}
+			}
+			cfg := newTestConfig(repo, tc.plan)
+			runner := setupTestRunner(t, cfg, tc.fake)
+
+			task := NewOneoffBackupTask(repo, tc.plan, time.Now(), tc.dryRun)
+			st := nextAndCreate(t, task, runner)
+
+			err := task.Run(context.Background(), st, runner)
+			if tc.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+
+			for _, cond := range tc.wantHooks {
+				assert.True(t, hookContains(runner.hookCalls, cond), "expected hook %v", cond)
+			}
+			for _, cond := range tc.wantNotHooks {
+				assert.False(t, hookContains(runner.hookCalls, cond), "unexpected hook %v", cond)
+			}
+
+			var scheduledTypes []string
+			for _, s := range runner.scheduledTasks {
+				scheduledTypes = append(scheduledTypes, s.Task.Type())
+			}
+			if tc.wantScheduled != nil {
+				assert.Equal(t, tc.wantScheduled, scheduledTypes)
+			}
+		})
+	}
+}
+
+// --- ForgetSnapshot task tests ---
+
+func TestForgetSnapshotTaskRun(t *testing.T) {
+	tests := []struct {
+		name    string
+		fake    *fakeRepoOrchestrator
+		wantErr bool
+	}{
+		{
+			name: "success",
+			fake: &fakeRepoOrchestrator{},
+		},
+		{
+			name:    "forget snapshot error",
+			fake:    &fakeRepoOrchestrator{forgetSnapshotErr: fmt.Errorf("forget failed")},
+			wantErr: true,
+		},
+		{
+			name:    "unlock error",
+			fake:    &fakeRepoOrchestrator{unlockErr: fmt.Errorf("unlock failed")},
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			repo := &v1.Repo{Id: "repo1", Guid: "guid1"}
+			cfg := newTestConfig(repo)
+			runner := setupTestRunner(t, cfg, tc.fake)
+
+			task := NewOneoffForgetSnapshotTask(repo, "plan1", 1, time.Now(), testSnapshotID)
+			st := nextAndCreate(t, task, runner)
+
+			err := task.Run(context.Background(), st, runner)
+			if tc.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				// On success, the task schedules an index snapshots task
+				require.Len(t, runner.scheduledTasks, 1)
+				assert.Equal(t, "index_snapshots", runner.scheduledTasks[0].Task.Type())
+			}
+		})
+	}
+}
+
+// --- ScheduledForgetTask tests ---
+
+func TestScheduledForgetTaskRun(t *testing.T) {
+	tests := []struct {
+		name          string
+		fake          *fakeRepoOrchestrator
+		wantErr       bool
+		wantHooks     []v1.Hook_Condition
+		wantScheduled int
+	}{
+		{
+			name:          "success",
+			fake:          &fakeRepoOrchestrator{},
+			wantHooks:     []v1.Hook_Condition{v1.Hook_CONDITION_FORGET_START, v1.Hook_CONDITION_FORGET_SUCCESS},
+			wantScheduled: 1, // stats task
+		},
+		{
+			name:    "forget error",
+			fake:    &fakeRepoOrchestrator{forgetErr: fmt.Errorf("forget failed")},
+			wantErr: true,
+			wantHooks: []v1.Hook_Condition{
+				v1.Hook_CONDITION_FORGET_START,
+				v1.Hook_CONDITION_FORGET_ERROR,
+				v1.Hook_CONDITION_ANY_ERROR,
+			},
+		},
+		{
+			name:    "unlock error",
+			fake:    &fakeRepoOrchestrator{unlockErr: fmt.Errorf("unlock failed")},
+			wantErr: true,
+			wantHooks: []v1.Hook_Condition{
+				v1.Hook_CONDITION_FORGET_ERROR,
+				v1.Hook_CONDITION_ANY_ERROR,
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			repo := &v1.Repo{
+				Id: "repo1", Guid: "guid1",
+				ForgetPolicy: &v1.ForgetPolicy{
+					Retention: &v1.RetentionPolicy{
+						Policy: &v1.RetentionPolicy_PolicyKeepLastN{PolicyKeepLastN: 5},
+					},
+				},
+			}
+			cfg := newTestConfig(repo)
+			runner := setupTestRunner(t, cfg, tc.fake)
+
+			task := NewScheduledForgetTask(repo, PlanForSystemTasks, true)
+			st := nextAndCreate(t, task, runner)
+
+			err := task.Run(context.Background(), st, runner)
+			if tc.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+
+			for _, cond := range tc.wantHooks {
+				assert.True(t, hookContains(runner.hookCalls, cond), "expected hook %v", cond)
+			}
+
+			assert.Len(t, runner.scheduledTasks, tc.wantScheduled)
+		})
+	}
+}
+
+// --- RestoreTask tests ---
+
+func TestRestoreTaskRun(t *testing.T) {
+	tests := []struct {
+		name    string
+		fake    *fakeRepoOrchestrator
+		wantErr bool
+	}{
+		{
+			name: "success",
+			fake: &fakeRepoOrchestrator{
+				restoreResult: &v1.RestoreProgressEntry{
+					MessageType:    "summary",
+					TotalFiles:     10,
+					TotalBytes:     5000,
+					FilesRestored:  10,
+					BytesRestored:  5000,
+				},
+			},
+		},
+		{
+			name:    "restore error",
+			fake:    &fakeRepoOrchestrator{restoreErr: fmt.Errorf("restore failed")},
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			repo := &v1.Repo{Id: "repo1", Guid: "guid1"}
+			cfg := newTestConfig(repo)
+			runner := setupTestRunner(t, cfg, tc.fake)
+
+			task := NewOneoffRestoreTask(repo, "plan1", 1, time.Now(), testSnapshotID, "/data", "/tmp/restore")
+			st := nextAndCreate(t, task, runner)
+
+			err := task.Run(context.Background(), st, runner)
+			if tc.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				restoreOp := st.Op.GetOperationRestore()
+				require.NotNil(t, restoreOp)
+				assert.NotNil(t, restoreOp.LastStatus)
+			}
+		})
+	}
+}
+
+// --- RunCommand tests ---
+
+func TestRunCommandTaskRun(t *testing.T) {
+	tests := []struct {
+		name    string
+		fake    *fakeRepoOrchestrator
+		wantErr bool
+	}{
+		{
+			name: "success",
+			fake: &fakeRepoOrchestrator{},
+		},
+		{
+			name:    "command error",
+			fake:    &fakeRepoOrchestrator{runCommandErr: fmt.Errorf("command failed")},
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			repo := &v1.Repo{Id: "repo1", Guid: "guid1"}
+			cfg := newTestConfig(repo)
+			runner := setupTestRunner(t, cfg, tc.fake)
+
+			task := NewOneoffRunCommandTask(repo, "plan1", 1, time.Now(), "echo hello")
+			st := nextAndCreate(t, task, runner)
+
+			err := task.Run(context.Background(), st, runner)
+			if tc.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}
+
+// --- IndexSnapshots tests ---
+
+func TestIndexSnapshotsTaskRun(t *testing.T) {
+	tests := []struct {
+		name    string
+		fake    *fakeRepoOrchestrator
+		wantErr bool
+	}{
+		{
+			name: "no snapshots",
+			fake: &fakeRepoOrchestrator{
+				snapshots: []*restic.Snapshot{},
+			},
+		},
+		{
+			name: "indexes new snapshots",
+			fake: &fakeRepoOrchestrator{
+				snapshots: []*restic.Snapshot{
+					{
+						Id:              testSnapshotID,
+						Time:            time.Now().Format(time.RFC3339Nano),
+						Tags:            []string{"plan:plan1", "created-by:test-instance"},
+						SnapshotSummary: restic.SnapshotSummary{},
+					},
+				},
+			},
+		},
+		{
+			name:    "snapshots error",
+			fake:    &fakeRepoOrchestrator{snapshotsErr: fmt.Errorf("snapshots failed")},
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			repo := &v1.Repo{Id: "repo1", Guid: "guid1"}
+			plan := &v1.Plan{Id: "plan1", Repo: "repo1"}
+			cfg := newTestConfig(repo, plan)
+			runner := setupTestRunner(t, cfg, tc.fake)
+
+			task := NewOneoffIndexSnapshotsTask(repo, time.Now())
+			st := nextAndCreate(t, task, runner)
+
+			err := task.Run(context.Background(), st, runner)
+			if tc.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}

internal/orchestrator/tasks/taskscheduledforget.go

@@ -1,231 +0,0 @@
-package tasks
-
-import (
-	"context"
-	"errors"
-	"fmt"
-	"time"
-
-	v1 "github.com/garethgeorge/backrest/gen/go/v1"
-	"github.com/garethgeorge/backrest/internal/oplog"
-	"github.com/garethgeorge/backrest/internal/protoutil"
-	"github.com/hashicorp/go-multierror"
-	"go.uber.org/zap"
-)
-
-type ScheduledForgetTask struct {
-	BaseTask
-	force  bool
-	didRun bool
-}
-
-func NewScheduledForgetTask(repo *v1.Repo, planID string, force bool) Task {
-	return &ScheduledForgetTask{
-		BaseTask: BaseTask{
-			TaskType:   "scheduled_forget",
-			TaskName:   fmt.Sprintf("scheduled forget for repo %q", repo.Id),
-			TaskRepo:   repo,
-			TaskPlanID: planID,
-		},
-		force: force,
-	}
-}
-
-func (t *ScheduledForgetTask) Next(now time.Time, runner TaskRunner) (ScheduledTask, error) {
-	if t.force {
-		if t.didRun {
-			return NeverScheduledTask, nil
-		}
-		t.didRun = true
-		return ScheduledTask{
-			Task:  t,
-			RunAt: now,
-			Op: &v1.Operation{
-				Op: &v1.Operation_OperationForget{},
-			},
-		}, nil
-	}
-
-	repo, err := runner.GetRepo(t.RepoID())
-	if err != nil {
-		return ScheduledTask{}, fmt.Errorf("get repo %v: %w", t.RepoID(), err)
-	}
-
-	if repo.GetForgetPolicy().GetSchedule() == nil {
-		return NeverScheduledTask, nil
-	}
-
-	var lastRan time.Time
-	var foundBackup bool
-	if err := runner.QueryOperations(oplog.Query{}.
-		SetInstanceID(runner.InstanceID()).
-		SetRepoGUID(repo.GetGuid()).
-		SetPlanID(PlanForSystemTasks).
-		SetReversed(true), func(op *v1.Operation) error {
-		if op.Status == v1.OperationStatus_STATUS_PENDING || op.Status == v1.OperationStatus_STATUS_SYSTEM_CANCELLED {
-			return nil
-		}
-		if _, ok := op.Op.(*v1.Operation_OperationForget); ok && op.UnixTimeEndMs != 0 {
-			lastRan = time.Unix(0, op.UnixTimeEndMs*int64(time.Millisecond))
-			return oplog.ErrStopIteration
-		}
-		if _, ok := op.Op.(*v1.Operation_OperationBackup); ok {
-			foundBackup = true
-		}
-		return nil
-	}); err != nil {
-		return NeverScheduledTask, fmt.Errorf("finding last scheduled forget run time: %w", err)
-	} else if !foundBackup {
-		lastRan = now
-	}
-
-	runAt, err := protoutil.ResolveSchedule(repo.GetForgetPolicy().GetSchedule(), lastRan, now)
-	if errors.Is(err, protoutil.ErrScheduleDisabled) {
-		return NeverScheduledTask, nil
-	} else if err != nil {
-		return NeverScheduledTask, fmt.Errorf("resolve schedule: %w", err)
-	}
-
-	return ScheduledTask{
-		Task:  t,
-		RunAt: runAt,
-		Op: &v1.Operation{
-			Op: &v1.Operation_OperationForget{},
-		},
-	}, nil
-}
-
-// shouldSkip returns true if there are no new successful backups since the last scheduled forget.
-func (t *ScheduledForgetTask) shouldSkip(runner TaskRunner, repo *v1.Repo) bool {
-	var lastForgetEndMs int64
-	var hasNewBackup bool
-
-	_ = runner.QueryOperations(oplog.Query{}.
-		SetInstanceID(runner.InstanceID()).
-		SetRepoGUID(repo.GetGuid()).
-		SetPlanID(PlanForSystemTasks).
-		SetReversed(true), func(op *v1.Operation) error {
-		if op.Status != v1.OperationStatus_STATUS_SUCCESS {
-			return nil
-		}
-		if _, ok := op.Op.(*v1.Operation_OperationForget); ok && op.UnixTimeEndMs != 0 {
-			lastForgetEndMs = op.UnixTimeEndMs
-			return oplog.ErrStopIteration
-		}
-		return nil
-	})
-
-	if lastForgetEndMs == 0 {
-		return false // no previous forget, don't skip
-	}
-
-	// Check if any backup completed after the last forget
-	_ = runner.QueryOperations(oplog.Query{}.
-		SetRepoGUID(repo.GetGuid()).
-		SetReversed(true), func(op *v1.Operation) error {
-		if op.UnixTimeEndMs < lastForgetEndMs {
-			return oplog.ErrStopIteration // older than last forget, stop looking
-		}
-		if op.Status == v1.OperationStatus_STATUS_SUCCESS {
-			if _, ok := op.Op.(*v1.Operation_OperationBackup); ok {
-				hasNewBackup = true
-				return oplog.ErrStopIteration
-			}
-		}
-		return nil
-	})
-
-	return !hasNewBackup
-}
-
-func (t *ScheduledForgetTask) Run(ctx context.Context, st ScheduledTask, runner TaskRunner) error {
-	op := st.Op
-
-	notifyError := func(err error) error {
-		return NotifyError(ctx, runner, t.Name(), err, v1.Hook_CONDITION_FORGET_ERROR)
-	}
-
-	repo, err := runner.GetRepo(t.RepoID())
-	if err != nil {
-		return notifyError(fmt.Errorf("get repo %q: %w", t.RepoID(), err))
-	}
-
-	// Skip if no new backups since last forget run.
-	// Mark as system-cancelled so it doesn't count as a successful run
-	// and the next schedule is computed from the last actual forget.
-	if t.shouldSkip(runner, repo) {
-		op.Op = &v1.Operation_OperationForget{
-			OperationForget: &v1.OperationForget{},
-		}
-		op.Status = v1.OperationStatus_STATUS_SYSTEM_CANCELLED
-		op.DisplayMessage = "Skipped: no new backups since last forget"
-		return nil
-	}
-
-	r, err := runner.GetRepoOrchestrator(t.RepoID())
-	if err != nil {
-		return notifyError(fmt.Errorf("get repo orchestrator %q: %w", t.RepoID(), err))
-	}
-
-	if err := r.UnlockIfAutoEnabled(ctx); err != nil {
-		return notifyError(fmt.Errorf("auto unlock repo %q: %w", t.RepoID(), err))
-	}
-
-	if err := runner.ExecuteHooks(ctx, []v1.Hook_Condition{
-		v1.Hook_CONDITION_FORGET_START,
-	}, HookVars{}); err != nil {
-		return notifyError(fmt.Errorf("forget start hook: %w", err))
-	}
-
-	forgot, err := r.ForgetAll(ctx, repo.GetForgetPolicy().GetRetention())
-
-	forgetOp := &v1.Operation_OperationForget{
-		OperationForget: &v1.OperationForget{
-			Forget: forgot,
-			Policy: repo.GetForgetPolicy().GetRetention(),
-		},
-	}
-	op.Op = forgetOp
-
-	// Mark forgotten snapshots in the oplog
-	var ops []*v1.Operation
-	for _, f := range forgot {
-		if e := runner.QueryOperations(oplog.Query{}.
-			SetRepoGUID(t.Repo().GetGuid()).
-			SetSnapshotID(f.Id), func(op *v1.Operation) error {
-			ops = append(ops, op)
-			return nil
-		}); e != nil {
-			err = multierror.Append(err, fmt.Errorf("lookup snapshot %v: %w", f.Id, e))
-		}
-	}
-
-	l := runner.Logger(ctx)
-	l.Sugar().Debugf("found %v snapshots were forgotten, marking this in oplog", len(ops))
-
-	for _, op := range ops {
-		if indexOp, ok := op.Op.(*v1.Operation_OperationIndexSnapshot); ok {
-			indexOp.OperationIndexSnapshot.Forgot = true
-			if e := runner.UpdateOperation(op); e != nil {
-				err = multierror.Append(err, fmt.Errorf("mark index snapshot %v as forgotten: %w", op.Id, e))
-			}
-		}
-	}
-
-	if err != nil {
-		return notifyError(fmt.Errorf("scheduled forget: %w", err))
-	}
-
-	// Schedule a stats task after successful forget
-	if e := runner.ScheduleTask(NewStatsTask(t.Repo(), PlanForSystemTasks, false), TaskPriorityStats); e != nil {
-		zap.L().Error("schedule stats task", zap.Error(e))
-	}
-
-	if e := runner.ExecuteHooks(ctx, []v1.Hook_Condition{
-		v1.Hook_CONDITION_FORGET_SUCCESS,
-	}, HookVars{}); e != nil {
-		return fmt.Errorf("forget end hook: %w", e)
-	}
-
-	return nil
-}

internal/orchestrator/tasks/testhelpers_test.go

@@ -0,0 +1,91 @@
+package tasks
+
+import (
+	"context"
+	"io"
+
+	v1 "github.com/garethgeorge/backrest/gen/go/v1"
+	"github.com/garethgeorge/backrest/pkg/restic"
+)
+
+// fakeRepoOrchestrator is a test double for the RepoOrchestrator interface.
+// Each method returns the corresponding configured result/error fields.
+type fakeRepoOrchestrator struct {
+	unlockErr error
+
+	backupResult *restic.BackupProgressEntry
+	backupErr    error
+
+	forgetResult []*v1.ResticSnapshot
+	forgetErr    error
+
+	forgetSnapshotErr error
+
+	pruneErr error
+	checkErr error
+
+	statsResult *v1.RepoStats
+	statsErr    error
+
+	restoreResult *v1.RestoreProgressEntry
+	restoreErr    error
+
+	snapshots    []*restic.Snapshot
+	snapshotsErr error
+
+	addTagsErr error
+
+	runCommandErr error
+}
+
+var _ RepoOrchestrator = &fakeRepoOrchestrator{}
+
+func (f *fakeRepoOrchestrator) UnlockIfAutoEnabled(ctx context.Context) error {
+	return f.unlockErr
+}
+
+func (f *fakeRepoOrchestrator) Backup(ctx context.Context, plan *v1.Plan, dryRun bool, cb func(event *restic.BackupProgressEntry)) (*restic.BackupProgressEntry, error) {
+	if cb != nil && f.backupResult != nil {
+		cb(f.backupResult)
+	}
+	return f.backupResult, f.backupErr
+}
+
+func (f *fakeRepoOrchestrator) Forget(ctx context.Context, policy *v1.RetentionPolicy, opts ...restic.GenericOption) ([]*v1.ResticSnapshot, error) {
+	return f.forgetResult, f.forgetErr
+}
+
+func (f *fakeRepoOrchestrator) ForgetSnapshot(ctx context.Context, snapshotId string) error {
+	return f.forgetSnapshotErr
+}
+
+func (f *fakeRepoOrchestrator) Prune(ctx context.Context, output io.Writer) error {
+	return f.pruneErr
+}
+
+func (f *fakeRepoOrchestrator) Check(ctx context.Context, output io.Writer) error {
+	return f.checkErr
+}
+
+func (f *fakeRepoOrchestrator) Stats(ctx context.Context) (*v1.RepoStats, error) {
+	return f.statsResult, f.statsErr
+}
+
+func (f *fakeRepoOrchestrator) Restore(ctx context.Context, snapshotId string, snapshotPath string, target string, cb func(event *v1.RestoreProgressEntry)) (*v1.RestoreProgressEntry, error) {
+	if cb != nil && f.restoreResult != nil {
+		cb(f.restoreResult)
+	}
+	return f.restoreResult, f.restoreErr
+}
+
+func (f *fakeRepoOrchestrator) Snapshots(ctx context.Context) ([]*restic.Snapshot, error) {
+	return f.snapshots, f.snapshotsErr
+}
+
+func (f *fakeRepoOrchestrator) AddTags(ctx context.Context, snapshotIDs []string, tags []string) error {
+	return f.addTagsErr
+}
+
+func (f *fakeRepoOrchestrator) RunCommand(ctx context.Context, command string, writer io.Writer) error {
+	return f.runCommandErr
+}

pkg/restic/restic.go

@@ -360,37 +360,26 @@ func (r *Repo) Forget(ctx context.Context, policy *RetentionPolicy, opts ...Gene
 		return nil, err
 	}
 
-	if len(results) != 1 {
-		return nil, fmt.Errorf("expected 1 output from forget, got %v", len(results))
+	if len(results) == 0 {
+		return nil, fmt.Errorf("expected at least 1 output from forget, got 0")
 	}
 
-	if err := results[0].Validate(); err != nil {
-		return nil, fmt.Errorf("invalid forget result: %w", err)
-	}
-
-	return &results[0], nil
-}
-
-// ForgetAll runs forget with a retention policy and returns results from all groups.
-// This is useful when running with --group-by tag which may return multiple groups.
-func (r *Repo) ForgetAll(ctx context.Context, policy *RetentionPolicy, opts ...GenericOption) ([]ForgetResult, error) {
-	args := []string{"forget", "--json"}
-	args = append(args, policy.toForgetFlags()...)
-
-	var results []ForgetResult
-	if err := r.executeWithJSONOutput(ctx, args, &results, opts...); err != nil {
-		return nil, err
-	}
-
-	for _, result := range results {
-		if err := result.Validate(); err != nil {
+	// Merge all groups into a single result. Restic returns one ForgetResult
+	// per group (e.g. when using --group-by tags), each with independent
+	// keep/remove lists.
+	merged := &ForgetResult{}
+	for _, r := range results {
+		if err := r.Validate(); err != nil {
 			return nil, fmt.Errorf("invalid forget result: %w", err)
 		}
+		merged.Keep = append(merged.Keep, r.Keep...)
+		merged.Remove = append(merged.Remove, r.Remove...)
 	}
 
-	return results, nil
+	return merged, nil
 }
 
+
 func (r *Repo) ForgetSnapshot(ctx context.Context, snapshotId string, opts ...GenericOption) error {
 	args := []string{"forget", "--json", snapshotId}
 	cmd := r.commandWithContext(ctx, args, opts...)

pkg/restic/restic_test.go

@@ -436,6 +436,65 @@ func TestResticForget(t *testing.T) {
 	}
 }
 
+func TestResticForgetMultiGroup(t *testing.T) {
+	t.Parallel()
+
+	repoDir := t.TempDir()
+	r := NewRepo(helpers.ResticBinary(t), repoDir, WithFlags("--no-cache"), WithEnv("RESTIC_PASSWORD=test"))
+	if err := r.Init(context.Background()); err != nil {
+		t.Fatalf("failed to init repo: %v", err)
+	}
+
+	testData := helpers.CreateTestData(t)
+
+	// Create snapshots with two different tags to produce multiple groups
+	var groupAIDs, groupBIDs []string
+	for i := 0; i < 5; i++ {
+		output, err := r.Backup(context.Background(), []string{testData}, nil, WithTags("group-a"))
+		if err != nil {
+			t.Fatalf("failed to backup group-a snapshot %d: %v", i, err)
+		}
+		groupAIDs = append(groupAIDs, output.SnapshotId)
+	}
+	for i := 0; i < 4; i++ {
+		output, err := r.Backup(context.Background(), []string{testData}, nil, WithTags("group-b"))
+		if err != nil {
+			t.Fatalf("failed to backup group-b snapshot %d: %v", i, err)
+		}
+		groupBIDs = append(groupBIDs, output.SnapshotId)
+	}
+
+	// Forget with --group-by tags, keeping 2 per group
+	res, err := r.Forget(context.Background(), &RetentionPolicy{KeepLastN: 2}, WithFlags("--group-by", "tags"))
+	if err != nil {
+		t.Fatalf("failed to forget snapshots: %v", err)
+	}
+
+	// Should keep 2 from each group = 4 total kept, 5 total removed (3 from A + 2 from B)
+	if len(res.Keep) != 4 {
+		t.Errorf("wanted 4 kept snapshots (2 per group), got: %d", len(res.Keep))
+	}
+	if len(res.Remove) != 5 {
+		t.Errorf("wanted 5 removed snapshots (3 from A + 2 from B), got: %d", len(res.Remove))
+	}
+
+	// Verify the kept snapshots are the most recent from each group
+	keptIDs := make(map[string]bool)
+	for _, s := range res.Keep {
+		keptIDs[s.Id] = true
+	}
+	for _, id := range groupAIDs[3:] {
+		if !keptIDs[id] {
+			t.Errorf("expected group-a snapshot %v to be kept", id)
+		}
+	}
+	for _, id := range groupBIDs[2:] {
+		if !keptIDs[id] {
+			t.Errorf("expected group-b snapshot %v to be kept", id)
+		}
+	}
+}
+
 func TestForgetSnapshotId(t *testing.T) {
 	t.Parallel()