fix: close Pi management final review findings

tools/thothctl/internal/pi/update.go

@@ -37,12 +37,13 @@ const (
 
 // Request contains only non-secret operator inputs.
 type Request struct {
-	StatePath string
-	Version   string
-	Source    Source
-	Image     string
-	Confirm   bool
-	Drain     bool
+	StatePath        string
+	RestartStatePath string
+	Version          string
+	Source           Source
+	Image            string
+	Confirm          bool
+	Drain            bool
 }
 
 // Result summarizes the completed, failed, or recovered transaction without command output.
@@ -72,6 +73,12 @@ func updateWithHooks(ctx context.Context, runner Runner, request Request, hooks
 	if request.StatePath == "" {
 		return Result{}, errors.New("update state path is required")
 	}
+	if request.RestartStatePath == "" {
+		request.RestartStatePath = pairedRestartStatePath(request.StatePath)
+	}
+	if err := validateRestartStatePaths(request.RestartStatePath, request.StatePath); err != nil {
+		return Result{StatePath: request.StatePath}, err
+	}
 	lock, err := acquireLock(request.StatePath)
 	if err != nil {
 		return Result{StatePath: request.StatePath}, err
@@ -96,6 +103,9 @@ func updateWithHooks(ctx context.Context, runner Runner, request Request, hooks
 		}
 		request.Image = canonical
 	}
+	if err := prepareLifecycleMutation(request.RestartStatePath, hooks.removeFile); err != nil {
+		return Result{StatePath: request.StatePath}, err
+	}
 	if old, err := readState(request.StatePath); err == nil && stateNeedsRecovery(old) {
 		return Result{StatePath: request.StatePath}, ErrInterruptedUpdate
 	} else if err != nil && !errors.Is(err, os.ErrNotExist) {
@@ -254,11 +264,17 @@ func failPreparation(statePath, overridePath string, state State, cause error, h
 }
 
 // Rollback restores the image recorded in durable update state. It is safe for interrupted runs.
-func Rollback(ctx context.Context, runner Runner, statePath string, confirm bool) (result Result, retErr error) {
-	return rollbackWithHooks(ctx, runner, statePath, confirm, defaultLifecycleHooks)
+func Rollback(ctx context.Context, runner Runner, statePath, restartStatePath string, confirm bool) (result Result, retErr error) {
+	return rollbackWithHooks(ctx, runner, statePath, restartStatePath, confirm, defaultLifecycleHooks)
 }
 
-func rollbackWithHooks(ctx context.Context, runner Runner, statePath string, confirm bool, hooks lifecycleHooks) (result Result, retErr error) {
+func rollbackWithHooks(ctx context.Context, runner Runner, statePath, restartStatePath string, confirm bool, hooks lifecycleHooks) (result Result, retErr error) {
+	if restartStatePath == "" {
+		restartStatePath = pairedRestartStatePath(statePath)
+	}
+	if err := validateRestartStatePaths(restartStatePath, statePath); err != nil {
+		return Result{StatePath: statePath}, err
+	}
 	lock, err := acquireLock(statePath)
 	if err != nil {
 		return Result{StatePath: statePath}, err
@@ -267,6 +283,9 @@ func rollbackWithHooks(ctx context.Context, runner Runner, statePath string, con
 	if !confirm {
 		return Result{StatePath: statePath}, ErrConfirmationRequired
 	}
+	if err := prepareLifecycleMutation(restartStatePath, hooks.removeFile); err != nil {
+		return Result{StatePath: statePath}, err
+	}
 	maintenanceErr := ensureMaintenance(ctx, runner)
 	clearMaintenance := maintenanceErr == nil
 	defer func() {