Clean up failed Notarius extraction runs

internal/stage/extract.go

@@ -151,19 +151,25 @@ func (extractStage) Run(ctx context.Context, env *Env, m *manifest.Manifest) (re
 		InputPath: inputPath, OutputRoot: outputRoot, WorkingDirectory: workingDirectory,
 		ReceiptPath: receiptPath, LogPath: logPath, Timeout: timeout, References: referenceBindings,
 	}
+	stageSucceeded := false
+	bundleMoved := false
+	defer func() {
+		if stageSucceeded {
+			return
+		}
+		if bundleMoved {
+			if cleanupErr := fileops.RemoveAllUnderRoot(filepath.Dir(durableBundle), durableBundle); cleanupErr != nil {
+				resultErr = errors.Join(resultErr, fmt.Errorf("extract: remove failed durable Notarius bundle: %w", cleanupErr))
+			}
+		}
+		if cleanupErr := fileops.RemoveAllUnderRoot(filepath.Dir(outputRoot), outputRoot); cleanupErr != nil {
+			resultErr = errors.Join(resultErr, fmt.Errorf("extract: remove failed Notarius output root: %w", cleanupErr))
+		}
+	}()
 	adapterResult, err := env.Notarius.Run(ctx, request)
 	if err != nil {
 		return nil, fmt.Errorf("extract: run notarius: %w", err)
 	}
-	bundleInstalled := false
-	defer func() {
-		if bundleInstalled {
-			return
-		}
-		if cleanupErr := os.RemoveAll(outputRoot); cleanupErr != nil {
-			resultErr = errors.Join(resultErr, fmt.Errorf("extract: remove failed Notarius output root: %w", cleanupErr))
-		}
-	}()
 	if strings.TrimSpace(adapterResult.BundleRoot) == "" || strings.TrimSpace(adapterResult.Index.Path) == "" {
 		return nil, fmt.Errorf("extract: notarius result is missing bundle or index path")
 	}
@@ -214,7 +220,7 @@ func (extractStage) Run(ctx context.Context, env *Env, m *manifest.Manifest) (re
 	if err := os.Rename(adapterResult.BundleRoot, durableBundle); err != nil {
 		return nil, fmt.Errorf("extract: install Notarius bundle: %w", err)
 	}
-	bundleInstalled = true
+	bundleMoved = true
 	if err := os.Remove(outputRoot); err != nil && !os.IsNotExist(err) {
 		return nil, fmt.Errorf("extract: remove empty Notarius output root: %w", err)
 	}
@@ -312,11 +318,13 @@ func (extractStage) Run(ctx context.Context, env *Env, m *manifest.Manifest) (re
 		"diagnostics":          boundedDiagnosticMetadata(adapterResult.Diagnostics),
 		"validation_summaries": boundedValidationMetadata(adapterResult.Receipt.ValidationSummaries),
 	}
-	return &StageResult{
+	result = &StageResult{
 		Outputs:  outputs,
 		Logs:     []string{receiptPath, logPath},
 		Metadata: metadata,
-	}, nil
+	}
+	stageSucceeded = true
+	return result, nil
 }
 
 func normalizeNotariusBundle(root string) error {