Wait for the interrupted operation to clean up before exit (closes #159)

On SIGINT/SIGTERM the process could exit before the interrupted command cleanup defers ran, leaving decrypted data in the temp directory (the blob cache and the decrypted snapshot database).

RunApp now mirrors fx run sequence: start, block on app.Wait(), then app.Stop(), returning only after Stop completes. fx delivers both an OS interrupt and the finished operation Shutdowner.Shutdown() on one channel. Stop runs the OnStop hooks; the operation hook cancels the command and waits for its goroutine to return (bounded by shutdownTimeout) before exit. The old code returned as soon as app.Done fired, without Stop, so a real interrupt unwound to os.Exit while cleanup still ran. Restore loops check the context between chunks and blobs so the wait ends promptly. A cli test drives RunApp through the OnStop hook.

Model: opus-4-8

internal/cli/app.go

@@ -7,12 +7,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"os"
-	"os/signal"
 	"path/filepath"
 	"strings"
 	"sync"
-	"syscall"
 	"time"
 
 	"github.com/adrg/xdg"
@@ -161,64 +158,45 @@ func cleanStartupError(err error) error {
 	return &startupError{msg: msg}
 }
 
-// RunApp starts and stops the fx application within the given context.
-// It handles graceful shutdown on interrupt signals (SIGINT, SIGTERM) and
-// ensures the application stops cleanly. The function blocks until the
-// application completes or is interrupted. Returns an error if startup fails.
+// RunApp starts the fx application, blocks until it is asked to stop, and
+// then stops it. The app is asked to stop either by an OS interrupt
+// (SIGINT/SIGTERM — fx installs its own handler when app.Wait is called) or,
+// on normal completion, by the finished operation calling
+// Shutdowner.Shutdown(); both arrive on the app.Wait channel.
+//
+// Stopping runs the fx OnStop hooks, and RunApp does not return until Stop
+// returns. On an interrupt the operation's OnStop hook cancels the running
+// command and waits for it to unwind — removing its decrypted scratch files —
+// so the process cannot proceed to exit mid-cleanup (issue #159). Waiting for
+// Stop before returning is what makes that hook effective: routing the
+// interrupt through app.Stop and not returning until it completes is required,
+// because fx also fires the app.Wait channel on the signal, and an earlier
+// version returned on that alone — unwinding to os.Exit while the concurrent
+// cleanup still ran. The stop is bounded by shutdownTimeout. Returns an error
+// if startup fails.
 func RunApp(ctx context.Context, app *fx.App) error {
-	// Set up signal handling for graceful shutdown
-	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
-
-	// Create a context that will be cancelled on signal
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-
-	// Start the app
 	err := app.Start(ctx)
 	if err != nil {
 		return cleanStartupError(err)
 	}
 
-	// Handle shutdown
-	shutdownComplete := make(chan struct{})
+	// Block until an interrupt or the finished operation's
+	// Shutdowner.Shutdown() arrives, then stop the app in this goroutine so we
+	// return only after its OnStop hooks — including the operation's cleanup
+	// wait — have run. Detach the stop from ctx's cancellation but keep its
+	// values, and bound it by shutdownTimeout.
+	<-app.Wait()
 
-	go func() {
-		defer close(shutdownComplete)
+	shutdownCtx, cancel := context.WithTimeout(
+		context.WithoutCancel(ctx), shutdownTimeout)
+	defer cancel()
 
-		<-sigChan
-		log.Notice("Received interrupt signal, shutting down gracefully...")
-
-		// Create a timeout context for shutdown. The parent ctx is being
-		// cancelled, so detach from its cancellation but keep its values.
-		shutdownCtx, shutdownCancel := context.WithTimeout(
-			context.WithoutCancel(ctx), shutdownTimeout)
-		defer shutdownCancel()
-
-		err := app.Stop(shutdownCtx)
-		if err != nil {
-			log.Error("Error during shutdown", "error", err)
-		}
-	}()
-
-	// Wait for the signal handler to complete shutdown or the app to
-	// request shutdown.
-	select {
-	case <-shutdownComplete:
-		// Shutdown completed via signal
-		return nil
-	case <-ctx.Done():
-		// Context cancelled (shouldn't happen in normal operation)
-		err := app.Stop(context.WithoutCancel(ctx))
-		if err != nil {
-			log.Error("Error stopping app", "error", err)
-		}
-
-		return ctx.Err()
-	case <-app.Done():
-		// App finished running (e.g., backup completed)
-		return nil
+	err = app.Stop(shutdownCtx)
+	if err != nil {
+		log.Error("Error during shutdown", "error", err)
 	}
+
+	return nil
 }
 
 // errReported marks a failure the operation has already shown the user
@@ -238,7 +216,10 @@ var errReported = errors.New("operation failed")
 //
 // op runs in a goroutine so OnStart returns promptly and an interrupt
 // can still cancel through OnStop; when it finishes, success or failure,
-// it triggers shutdown, which is what lets RunWithApp return. report is
+// it triggers shutdown, which is what lets RunWithApp return. On an
+// interrupt OnStop cancels op and waits for the goroutine to return, so
+// op's cleanup (removing decrypted scratch files) runs before the
+// process exits; the wait is bounded by shutdownTimeout. report is
 // called with a non-canceled failure so the caller can log it (and
 // suppress it under --json) before it becomes errReported. A context
 // cancellation is the interrupt path, not a failure: it is neither
@@ -254,9 +235,11 @@ func RunOperation(
 
 	opts.Invokes = append(opts.Invokes,
 		fx.Invoke(func(v *vaultik.Vaultik, lc fx.Lifecycle) {
+			var stop func(context.Context) bool
+
 			lc.Append(fx.Hook{
 				OnStart: func(_ context.Context) error {
-					go func() {
+					stop = v.StartOperation(func() {
 						err := op(v)
 						if err != nil && !errors.Is(err, context.Canceled) {
 							report(err)
@@ -270,12 +253,20 @@ func RunOperation(
 						if stopErr != nil {
 							log.Error("Failed to shutdown", "error", stopErr)
 						}
-					}()
+					})
 
 					return nil
 				},
-				OnStop: func(_ context.Context) error {
-					v.Cancel()
+				// On an interrupt, cancel the operation and wait for it to
+				// unwind so its cleanup defers (which remove decrypted
+				// scratch files from the temp directory) run before the
+				// process exits. The wait is bounded by ctx, the existing
+				// shutdownTimeout.
+				OnStop: func(ctx context.Context) error {
+					if !stop(ctx) {
+						log.Warn("Shutdown timed out before the operation " +
+							"finished; decrypted temporary files may remain")
+					}
 
 					return nil
 				},

internal/cli/runapp_interrupt_test.go

@@ -0,0 +1,97 @@
+package cli_test
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/fx"
+	"sneak.berlin/go/vaultik/internal/cli"
+)
+
+// TestRunAppWaitsForOperationCleanupOnShutdown drives RunApp with an fx app
+// wired the way RunOperation wires a command: a single lifecycle hook whose
+// OnStart launches the operation in its own goroutine and whose OnStop cancels
+// it and blocks until that goroutine returns. The operation stands in for a
+// restore blocked mid-download — it holds a decrypted "scratch" file and only
+// removes it as it unwinds on cancellation.
+//
+// The app is asked to stop once the operation is running (standing in for an
+// OS interrupt; fx delivers a real signal and Shutdowner.Shutdown() on the
+// same app.Wait channel, so both drive the identical shutdown path). RunApp
+// must not return until app.Stop has run the OnStop hook, so the scratch file
+// must be gone by the time RunApp returns. Before the fix RunApp returned as
+// soon as the app.Wait/Done channel fired, without running app.Stop, so the
+// cleanup never ran and this file would still be on disk (issue #159).
+func TestRunAppWaitsForOperationCleanupOnShutdown(t *testing.T) {
+	t.Parallel()
+
+	scratch := filepath.Join(t.TempDir(), "decrypted-scratch")
+	require.NoError(t, os.WriteFile(scratch, []byte("secret"), 0o600))
+
+	// Cancel and reap the operation even if RunApp returns without doing so
+	// (the buggy path), so the goroutine cannot leak past the test.
+	opCtx, opCancel := context.WithCancel(context.Background())
+	t.Cleanup(opCancel)
+
+	var stop func(context.Context) bool
+
+	app := fx.New(
+		fx.NopLogger,
+		fx.Invoke(func(lc fx.Lifecycle, sh fx.Shutdowner) {
+			lc.Append(fx.Hook{
+				OnStart: func(_ context.Context) error {
+					done := make(chan struct{})
+
+					go func() {
+						defer close(done)
+
+						// Blocked mid-operation until cancelled, then run the
+						// cleanup an interrupted restore would run.
+						<-opCtx.Done()
+
+						_ = os.Remove(scratch)
+					}()
+
+					stop = func(ctx context.Context) bool {
+						opCancel()
+
+						select {
+						case <-done:
+							return true
+						case <-ctx.Done():
+							return false
+						}
+					}
+
+					// Ask the app to stop now that the operation is running.
+					go func() { _ = sh.Shutdown() }()
+
+					return nil
+				},
+				OnStop: func(ctx context.Context) error {
+					stop(ctx)
+
+					return nil
+				},
+			})
+		}),
+	)
+
+	done := make(chan error, 1)
+	go func() { done <- cli.RunApp(context.Background(), app) }()
+
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(30 * time.Second):
+		t.Fatal("RunApp did not return after shutdown was requested")
+	}
+
+	_, err := os.Stat(scratch)
+	require.True(t, os.IsNotExist(err),
+		"RunApp returned before the operation removed its decrypted scratch file")
+}