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 {

internal/stage/extract_test.go

@@ -492,11 +492,31 @@ func TestExtractStageValidatesSelectedLaneJSON(t *testing.T) {
 func TestExtractStageAdapterAndPromotionFailuresReturnNoOutputs(t *testing.T) {
 	t.Run("adapter", func(t *testing.T) {
 		env, m, fake := setupExtractEnv(t)
+		fixture := extractFixtureFromEnv(t, env, m)
+		receiptPath := artifacts.SessionRunNotariusReceiptPathForCampaign(fixture.workspace, fixture.campaign, fixture.sessionID, fixture.runID)
+		logPath := artifacts.SessionRunNotariusLogPathForCampaign(fixture.workspace, fixture.campaign, fixture.sessionID, fixture.runID)
+		for _, path := range []string{receiptPath, logPath} {
+			if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
+				t.Fatalf("MkdirAll(diagnostic parent) error = %v", err)
+			}
+			if err := os.WriteFile(path, []byte("diagnostic"), 0o644); err != nil {
+				t.Fatalf("WriteFile(diagnostic) error = %v", err)
+			}
+		}
 		fake.Err = errors.New("adapter failed")
 		result, err := (extractStage{}).Run(context.Background(), env, m)
 		if err == nil || !strings.Contains(err.Error(), "adapter failed") || result != nil {
 			t.Fatalf("Run() result = %#v error = %v", result, err)
 		}
+		outputRoot := artifacts.SessionRunNotariusOutputRootForCampaign(fixture.workspace, fixture.campaign, fixture.sessionID, fixture.runID)
+		if _, err := os.Stat(outputRoot); !os.IsNotExist(err) {
+			t.Fatalf("failed adapter output root still exists: %v", err)
+		}
+		for _, path := range []string{receiptPath, logPath} {
+			if _, err := os.Stat(path); err != nil {
+				t.Fatalf("diagnostic %q was not retained: %v", path, err)
+			}
+		}
 	})
 
 	t.Run("promotion", func(t *testing.T) {
@@ -514,6 +534,27 @@ func TestExtractStageAdapterAndPromotionFailuresReturnNoOutputs(t *testing.T) {
 			t.Fatalf("failed installation staging bundle still exists: %v", err)
 		}
 	})
+
+	t.Run("failure after rename", func(t *testing.T) {
+		env, m, _ := setupExtractEnv(t)
+		fixture := extractFixtureFromEnv(t, env, m)
+		outputRoot := artifacts.SessionRunNotariusOutputRootForCampaign(fixture.workspace, fixture.campaign, fixture.sessionID, fixture.runID)
+		if err := os.WriteFile(filepath.Join(outputRoot, "unexpected-sidecar"), []byte("unexpected"), 0o644); err != nil {
+			t.Fatalf("WriteFile(unexpected sidecar) error = %v", err)
+		}
+		destination := artifacts.SessionNotariusBundleDirForCampaign(fixture.workspace, fixture.campaign, fixture.sessionID, fixture.runID)
+
+		result, err := (extractStage{}).Run(context.Background(), env, m)
+		if err == nil || !strings.Contains(err.Error(), "remove empty Notarius output root") || result != nil {
+			t.Fatalf("Run() result = %#v error = %v", result, err)
+		}
+		if _, err := os.Stat(destination); !os.IsNotExist(err) {
+			t.Fatalf("failed stage durable bundle still exists: %v", err)
+		}
+		if _, err := os.Stat(outputRoot); !os.IsNotExist(err) {
+			t.Fatalf("failed stage output root still exists: %v", err)
+		}
+	})
 }
 
 func TestExtractionFingerprintIsIndependentOfOutputMapOrder(t *testing.T) {