From da3bd9e31c72e29a2a75ef82107d05e7f91cc4db Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Fri, 18 Sep 2026 04:08:34 -0700
Subject: [PATCH] test(ui): model E2E cleanup failures as values

---
 tests/e2e/ui/helpers/roundTrip.ts | 111 ++++++++++++++++++++++--------
 1 file changed, 81 insertions(+), 30 deletions(-)

diff --git a/tests/e2e/ui/helpers/roundTrip.ts b/tests/e2e/ui/helpers/roundTrip.ts
index 1fc2d0aec1a..91e48b7087c 100644
--- a/tests/e2e/ui/helpers/roundTrip.ts
+++ b/tests/e2e/ui/helpers/roundTrip.ts
@@ -33,39 +33,90 @@ export async function readBack<T = any>(
   return (await res.json()) as T;
 }
 
-export async function runWithCleanup(
-  action: () => Promise<void>,
-  cleanup: () => Promise<boolean>,
-): Promise<void> {
-  const outcome = await Promise.resolve()
+type OperationOutcome =
+  | { readonly status: "success" }
+  | { readonly status: "failure"; readonly error: unknown };
+
+type RunFailure =
+  | { readonly status: "action_failure"; readonly error: unknown }
+  | { readonly status: "cleanup_failure"; readonly error: unknown }
+  | {
+      readonly status: "action_and_cleanup_failure";
+      readonly actionError: unknown;
+      readonly cleanupError: unknown;
+    };
+
+function toRunFailure(
+  actionOutcome: OperationOutcome,
+  cleanupOutcome: OperationOutcome,
+): RunFailure | null {
+  if (
+    actionOutcome.status === "failure" &&
+    cleanupOutcome.status === "failure"
+  ) {
+    return {
+      status: "action_and_cleanup_failure",
+      actionError: actionOutcome.error,
+      cleanupError: cleanupOutcome.error,
+    };
+  }
+  if (actionOutcome.status === "failure") {
+    return { status: "action_failure", error: actionOutcome.error };
+  }
+  if (cleanupOutcome.status === "failure") {
+    return { status: "cleanup_failure", error: cleanupOutcome.error };
+  }
+  return null;
+}
+
+function raiseRunFailure(failure: RunFailure): never {
+  switch (failure.status) {
+    case "action_failure":
+      throw failure.error;
+    case "cleanup_failure":
+      throw failure.error;
+    case "action_and_cleanup_failure":
+      throw new AggregateError(
+        [failure.actionError, failure.cleanupError],
+        "Action and cleanup failed",
+      );
+  }
+}
+
+async function runAction(
+  action: () => void | Promise<void>,
+): Promise<OperationOutcome> {
+  return Promise.resolve()
     .then(action)
     .then(
       () => ({ status: "success" as const }),
       (error: unknown) => ({ status: "failure" as const, error }),
     );
-  try {
-    if (outcome.status === "failure") throw outcome.error;
-  } finally {
-    const cleanupOutcome = await Promise.resolve()
-      .then(cleanup)
-      .then(
-        (succeeded) =>
-          succeeded
-            ? { status: "success" as const }
-            : {
-                status: "failure" as const,
-                error: new Error("Failed to clean up UI E2E resource"),
-              },
-        (error: unknown) => ({ status: "failure" as const, error }),
-      );
-    if (cleanupOutcome.status === "failure") {
-      if (outcome.status === "failure") {
-        throw new AggregateError(
-          [outcome.error, cleanupOutcome.error],
-          "Action and cleanup failed",
-        );
-      }
-      throw cleanupOutcome.error;
-    }
-  }
+}
+
+async function runCleanup(
+  cleanup: () => boolean | Promise<boolean>,
+): Promise<OperationOutcome> {
+  return Promise.resolve()
+    .then(cleanup)
+    .then(
+      (succeeded) =>
+        succeeded
+          ? { status: "success" as const }
+          : {
+              status: "failure" as const,
+              error: new Error("Failed to clean up UI E2E resource"),
+            },
+      (error: unknown) => ({ status: "failure" as const, error }),
+    );
+}
+
+export async function runWithCleanup(
+  action: () => void | Promise<void>,
+  cleanup: () => boolean | Promise<boolean>,
+): Promise<void> {
+  const actionOutcome = await runAction(action);
+  const cleanupOutcome = await runCleanup(cleanup);
+  const failure = toRunFailure(actionOutcome, cleanupOutcome);
+  if (failure !== null) raiseRunFailure(failure);
 }