test(app): abort sessions and wait for idle before e2e cleanup (#16439)

3 files changed, 6 insertions, 5 deletions

diff --git a/packages/app/e2e/prompt/prompt-async.spec.ts b/packages/app/e2e/prompt/prompt-async.spec.ts
index 10e3fc312..51fbc3e4a 100644
--- a/packages/app/e2e/prompt/prompt-async.spec.ts
+++ b/packages/app/e2e/prompt/prompt-async.spec.ts
@@ -1,6 +1,6 @@
 import { test, expect } from "../fixtures"
 import { promptSelector } from "../selectors"
-import { sessionIDFromUrl, withSession } from "../actions"
+import { cleanupSession, sessionIDFromUrl, withSession } from "../actions"
 
 const text = (value: string | null) => (value ?? "").replace(/\u200B/g, "").trim()
 
@@ -40,7 +40,7 @@ test("prompt succeeds when sync message endpoint is unreachable", async ({ page,
       )
       .toContain(token)
   } finally {
-    await sdk.session.delete({ sessionID }).catch(() => undefined)
+    await cleanupSession({ sdk, sessionID })
   }
 })
 
diff --git a/packages/app/e2e/prompt/prompt-shell.spec.ts b/packages/app/e2e/prompt/prompt-shell.spec.ts
index c9880bf20..4c92f4a2f 100644
--- a/packages/app/e2e/prompt/prompt-shell.spec.ts
+++ b/packages/app/e2e/prompt/prompt-shell.spec.ts
@@ -14,7 +14,7 @@ const isBash = (part: unknown): part is ToolPart => {
 test("shell mode runs a command in the project directory", async ({ page, withProject }) => {
   test.setTimeout(120_000)
 
-  await withProject(async ({ directory, gotoSession }) => {
+  await withProject(async ({ directory, gotoSession, trackSession }) => {
     const sdk = createSdk(directory)
     const prompt = page.locator(promptSelector)
     const cmd = process.platform === "win32" ? "dir" : "ls"
@@ -31,6 +31,7 @@ test("shell mode runs a command in the project directory", async ({ page, withPr
 
     const id = sessionIDFromUrl(page.url())
     if (!id) throw new Error(`Failed to parse session id from url: ${page.url()}`)
+    trackSession(id, directory)
 
     await expect
       .poll(
diff --git a/packages/app/e2e/prompt/prompt.spec.ts b/packages/app/e2e/prompt/prompt.spec.ts
index ff9f5daf0..0466d0988 100644
--- a/packages/app/e2e/prompt/prompt.spec.ts
+++ b/packages/app/e2e/prompt/prompt.spec.ts
@@ -1,6 +1,6 @@
 import { test, expect } from "../fixtures"
 import { promptSelector } from "../selectors"
-import { sessionIDFromUrl, withSession } from "../actions"
+import { cleanupSession, sessionIDFromUrl, withSession } from "../actions"
 
 test("can send a prompt and receive a reply", async ({ page, sdk, gotoSession }) => {
   test.setTimeout(120_000)
@@ -46,7 +46,7 @@ test("can send a prompt and receive a reply", async ({ page, sdk, gotoSession })
       .toContain(token)
   } finally {
     page.off("pageerror", onPageError)
-    await sdk.session.delete({ sessionID }).catch(() => undefined)
+    await cleanupSession({ sdk, sessionID })
   }
 
   if (pageErrors.length > 0) {