From 60706feb348cd1c8b18c2218e5d0305535ab1233 Mon Sep 17 00:00:00 2001
From: Mohamed Boudra <boudra.moha@gmail.com>
Date: Tue, 18 Aug 2026 23:34:05 +0200
Subject: [PATCH] fix(test): stabilize tool call shimmer lifecycle gate in
 browser e2e

---
 .../app/e2e/browser/tool-call-shimmer.spec.ts | 45 +++++++------------
 1 file changed, 15 insertions(+), 30 deletions(-)

diff --git a/packages/app/e2e/browser/tool-call-shimmer.spec.ts b/packages/app/e2e/browser/tool-call-shimmer.spec.ts
index 58660c3df..71ef6555e 100644
--- a/packages/app/e2e/browser/tool-call-shimmer.spec.ts
+++ b/packages/app/e2e/browser/tool-call-shimmer.spec.ts
@@ -1,6 +1,7 @@
 import type { Locator, Page } from "@playwright/test";
 import { test, expect } from "../support/fixtures";
 import { daemonWsRoutePattern } from "../support/helpers/daemon-port";
+import { expectComposerVisible } from "../support/helpers/composer";
 import { openAgentRoute, seedMockAgentWorkspace } from "../support/helpers/mock-agent";
 
 type WebSocketMessage = string | Buffer;
@@ -64,9 +65,8 @@ async function gateSecondToolCall(page: Page, agentId: string) {
   let firstCallId: string | null = null;
   let secondCallId: string | null = null;
   let secondRunningMessage: WebSocketMessage | null = null;
+  let pauseAfterSecondRunning = false;
   let releaseSecondRequested = false;
-  let pauseServerMessages = false;
-  let secondRunningForwarded = false;
   let forwardSecondRunning: (() => void) | null = null;
   let resolveFirstCompleted!: () => void;
   let resolveSecondRunning!: () => void;
@@ -80,15 +80,14 @@ async function gateSecondToolCall(page: Page, agentId: string) {
   await page.routeWebSocket(daemonWsRoutePattern(), (ws) => {
     const server = ws.connectToServer();
     forwardSecondRunning = () => {
-      if (!secondRunningMessage || secondRunningForwarded) {
-        return;
+      if (secondRunningMessage) {
+        ws.send(secondRunningMessage);
+        secondRunningMessage = null;
       }
-      ws.send(secondRunningMessage);
-      secondRunningForwarded = true;
     };
     ws.onMessage((message) => server.send(message));
     server.onMessage((message) => {
-      if (pauseServerMessages) {
+      if (pauseAfterSecondRunning) {
         return;
       }
 
@@ -97,6 +96,7 @@ async function gateSecondToolCall(page: Page, agentId: string) {
         ws.send(message);
         return;
       }
+
       if (!firstCallId) {
         firstCallId = toolCall.callId;
       }
@@ -109,24 +109,11 @@ async function gateSecondToolCall(page: Page, agentId: string) {
       }
 
       secondCallId ??= toolCall.callId;
-      if (
-        toolCall.callId === secondCallId &&
-        (toolCall.status === "running" || toolCall.status === "executing")
-      ) {
-        secondRunningMessage = message;
-        pauseServerMessages = true;
-        resolveSecondRunning();
-        if (releaseSecondRequested) {
-          forwardSecondRunning?.();
-        }
-        return;
-      }
-
-      if (toolCall.callId === secondCallId && toolCall.status === "completed") {
-        // Keep the second real tool call inspectably active. Production providers
-        // send this same status shape while their work remains in flight.
-        secondRunningMessage = replaceToolCallStatus(message, "running");
-        pauseServerMessages = true;
+      if (toolCall.callId === secondCallId) {
+        const transformedMessage =
+          toolCall.status === "completed" ? replaceToolCallStatus(message, "running") : message;
+        secondRunningMessage = transformedMessage;
+        pauseAfterSecondRunning = true;
         resolveSecondRunning();
         if (releaseSecondRequested) {
           forwardSecondRunning?.();
@@ -143,10 +130,7 @@ async function gateSecondToolCall(page: Page, agentId: string) {
     waitForSecondRunning: () => secondRunning,
     releaseSecondRunning() {
       releaseSecondRequested = true;
-      pauseServerMessages = false;
-      if (secondRunningMessage) {
-        forwardSecondRunning?.();
-      }
+      forwardSecondRunning?.();
     },
   };
 }
@@ -192,10 +176,11 @@ test("measures an overview heading that becomes loading after its idle mount", a
       workspaceId: agent.workspaceId,
       agentId: agent.agentId,
     });
+    await expectComposerVisible(page);
     await agent.client.sendAgentMessage(agent.agentId, "Prove the overview shimmer lifecycle.");
 
     await gate.waitForFirstCompleted();
-    const group = page.getByTestId("tool-call-group");
+    const group = page.getByTestId("tool-call-group").first();
     await expect(group).toBeVisible();
     const idleGroupHandle = await group.elementHandle();
     if (!idleGroupHandle) {