fix(test): auto-acknowledge tool-result follow-ups in mock LLM server (#20528)

packages/app/e2e/session/session-review.spec.ts

@@ -1,6 +1,6 @@
 import { waitSessionIdle, withSession } from "../actions"
 import { test, expect } from "../fixtures"
-import { createSdk } from "../utils"
+import { inputMatch } from "../prompt/mock"
 
 const count = 14
 
@@ -40,7 +40,14 @@ function edit(file: string, prev: string, next: string) {
   )
 }
 
-async function patch(sdk: ReturnType<typeof createSdk>, sessionID: string, patchText: string) {
+async function patchWithMock(
+  llm: Parameters<typeof test>[0]["llm"],
+  sdk: Parameters<typeof withSession>[0],
+  sessionID: string,
+  patchText: string,
+) {
+  const callsBefore = await llm.calls()
+  await llm.toolMatch(inputMatch({ patchText }), "apply_patch", { patchText })
   await sdk.session.promptAsync({
     sessionID,
     agent: "build",
@@ -54,6 +61,13 @@ async function patch(sdk: ReturnType<typeof createSdk>, sessionID: string, patch
     parts: [{ type: "text", text: "Apply the provided patch exactly once." }],
   })
 
+  // Wait for the agent loop to actually start before checking idle.
+  // promptAsync is fire-and-forget — without this, waitSessionIdle can
+  // return immediately because the session status is still undefined.
+  await expect
+    .poll(() => llm.calls().then((c) => c > callsBefore), { timeout: 30_000 })
+    .toBe(true)
+
   await waitSessionIdle(sdk, sessionID, 120_000)
 }
 
@@ -233,8 +247,7 @@ async function fileOverflow(page: Parameters<typeof test>[0]["page"]) {
   }
 }
 
-test("review applies inline comment clicks without horizontal overflow", async ({ page, withProject }) => {
-  test.skip(true, "Flaky in CI for now.")
+test("review applies inline comment clicks without horizontal overflow", async ({ page, llm, withMockProject }) => {
   test.setTimeout(180_000)
 
   const tag = `review-comment-${Date.now()}`
@@ -243,16 +256,15 @@ test("review applies inline comment clicks without horizontal overflow", async (
 
   await page.setViewportSize({ width: 1280, height: 900 })
 
-  await withProject(async (project) => {
-    const sdk = createSdk(project.directory)
-
-    await withSession(sdk, `e2e review comment ${tag}`, async (session) => {
-      await patch(sdk, session.id, seed([{ file, mark: tag }]))
+  await withMockProject(async (project) => {
+    await withSession(project.sdk, `e2e review comment ${tag}`, async (session) => {
+      project.trackSession(session.id)
+      await patchWithMock(llm, project.sdk, session.id, seed([{ file, mark: tag }]))
 
       await expect
         .poll(
           async () => {
-            const diff = await sdk.session.diff({ sessionID: session.id }).then((res) => res.data ?? [])
+            const diff = await project.sdk.session.diff({ sessionID: session.id }).then((res) => res.data ?? [])
             return diff.length
           },
           { timeout: 60_000 },
@@ -283,8 +295,7 @@ test("review applies inline comment clicks without horizontal overflow", async (
   })
 })
 
-test("review file comments submit on click without clipping actions", async ({ page, withProject }) => {
-  test.skip(true, "Flaky in CI for now.")
+test("review file comments submit on click without clipping actions", async ({ page, llm, withMockProject }) => {
   test.setTimeout(180_000)
 
   const tag = `review-file-comment-${Date.now()}`
@@ -293,16 +304,15 @@ test("review file comments submit on click without clipping actions", async ({ p
 
   await page.setViewportSize({ width: 1280, height: 900 })
 
-  await withProject(async (project) => {
-    const sdk = createSdk(project.directory)
-
-    await withSession(sdk, `e2e review file comment ${tag}`, async (session) => {
-      await patch(sdk, session.id, seed([{ file, mark: tag }]))
+  await withMockProject(async (project) => {
+    await withSession(project.sdk, `e2e review file comment ${tag}`, async (session) => {
+      project.trackSession(session.id)
+      await patchWithMock(llm, project.sdk, session.id, seed([{ file, mark: tag }]))
 
       await expect
         .poll(
           async () => {
-            const diff = await sdk.session.diff({ sessionID: session.id }).then((res) => res.data ?? [])
+            const diff = await project.sdk.session.diff({ sessionID: session.id }).then((res) => res.data ?? [])
             return diff.length
           },
           { timeout: 60_000 },
@@ -334,8 +344,7 @@ test("review file comments submit on click without clipping actions", async ({ p
   })
 })
 
-test("review keeps scroll position after a live diff update", async ({ page, withProject }) => {
-  test.skip(Boolean(process.env.CI), "Flaky in CI for now.")
+test.fixme("review keeps scroll position after a live diff update", async ({ page, llm, withMockProject }) => {
   test.setTimeout(180_000)
 
   const tag = `review-${Date.now()}`
@@ -345,16 +354,15 @@ test("review keeps scroll position after a live diff update", async ({ page, wit
 
   await page.setViewportSize({ width: 1600, height: 1000 })
 
-  await withProject(async (project) => {
-    const sdk = createSdk(project.directory)
-
-    await withSession(sdk, `e2e review ${tag}`, async (session) => {
-      await patch(sdk, session.id, seed(list))
+  await withMockProject(async (project) => {
+    await withSession(project.sdk, `e2e review ${tag}`, async (session) => {
+      project.trackSession(session.id)
+      await patchWithMock(llm, project.sdk, session.id, seed(list))
 
       await expect
         .poll(
           async () => {
-            const info = await sdk.session.get({ sessionID: session.id }).then((res) => res.data)
+            const info = await project.sdk.session.get({ sessionID: session.id }).then((res) => res.data)
             return info?.summary?.files ?? 0
           },
           { timeout: 60_000 },
@@ -364,7 +372,7 @@ test("review keeps scroll position after a live diff update", async ({ page, wit
       await expect
         .poll(
           async () => {
-            const diff = await sdk.session.diff({ sessionID: session.id }).then((res) => res.data ?? [])
+            const diff = await project.sdk.session.diff({ sessionID: session.id }).then((res) => res.data ?? [])
             return diff.length
           },
           { timeout: 60_000 },
@@ -381,15 +389,16 @@ test("review keeps scroll position after a live diff update", async ({ page, wit
       const view = page.locator('[data-slot="session-review-scroll"] .scroll-view__viewport').first()
       await expect(view).toBeVisible()
       const heads = page.getByRole("heading", { level: 3 }).filter({ hasText: /^review-scroll-/ })
-      await expect(heads).toHaveCount(list.length, {
-        timeout: 60_000,
-      })
+      await expect(heads).toHaveCount(list.length, { timeout: 60_000 })
 
       await expand(page)
       await waitMark(page, hit.file, hit.mark)
 
       const row = page
-        .getByRole("heading", { level: 3, name: new RegExp(hit.file.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")) })
+        .getByRole("heading", {
+          level: 3,
+          name: new RegExp(hit.file.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")),
+        })
         .first()
       await expect(row).toBeVisible()
       await row.evaluate((el) => el.scrollIntoView({ block: "center" }))
@@ -398,12 +407,12 @@ test("review keeps scroll position after a live diff update", async ({ page, wit
       const prev = await spot(page, hit.file)
       if (!prev) throw new Error(`missing review row for ${hit.file}`)
 
-      await patch(sdk, session.id, edit(hit.file, hit.mark, next))
+      await patchWithMock(llm, project.sdk, session.id, edit(hit.file, hit.mark, next))
 
       await expect
         .poll(
           async () => {
-            const diff = await sdk.session.diff({ sessionID: session.id }).then((res) => res.data ?? [])
+            const diff = await project.sdk.session.diff({ sessionID: session.id }).then((res) => res.data ?? [])
             const item = diff.find((item) => item.file === hit.file)
             return typeof item?.after === "string" ? item.after : ""
           },