Fix flaky React and chat concurrency tests (#1426)

Made-with: Cursor

Author: threepointone

packages/agents/src/react-tests/setup.ts

@@ -32,6 +32,30 @@ function isPortAvailable(port: number): Promise<boolean> {
   });
 }
 
+async function isWorkerReachable(port: number): Promise<boolean> {
+  try {
+    const response = await fetch(`http://127.0.0.1:${port}`, {
+      signal: AbortSignal.timeout(500)
+    });
+    // Any HTTP response means a worker is already serving the test port.
+    return response.status >= 100;
+  } catch {
+    return false;
+  }
+}
+
+async function waitForReachableWorker(
+  port: number,
+  timeoutMs: number
+): Promise<boolean> {
+  const deadline = Date.now() + timeoutMs;
+  while (Date.now() < deadline) {
+    if (await isWorkerReachable(port)) return true;
+    await new Promise((r) => setTimeout(r, 100));
+  }
+  return false;
+}
+
 /**
  * Kill any process listening on the given port.
  * Handles stale processes left behind by previous test runs that were
@@ -77,6 +101,13 @@ export async function setup() {
   // that was forcefully terminated before teardown could run.
   const portAvailable = await isPortAvailable(TEST_WORKER_PORT);
   if (!portAvailable) {
+    if (await waitForReachableWorker(TEST_WORKER_PORT, 2000)) {
+      console.log(
+        `[setup] Reusing test worker at http://127.0.0.1:${TEST_WORKER_PORT}`
+      );
+      return;
+    }
+
     console.log(
       `[setup] Port ${TEST_WORKER_PORT} in use — killing stale process...`
     );

packages/agents/src/react-tests/useAgent.test.tsx

@@ -377,27 +377,41 @@ describe("useAgent hook", () => {
 
     it("should update state property on server broadcast", async () => {
       const { host, protocol } = getTestWorkerHost();
-      let capturedAgent: TestAgent | null = null;
+      let observerAgent: TestAgent | null = null;
+      let senderAgent: TestAgent | null = null;
+      const room = `hook-test-state-prop-server-${crypto.randomUUID()}`;
 
       const { container } = await render(
         <SuspenseWrapper>
           <StateTrackingComponent
             options={{
               agent: "TestStateAgent",
-              name: "hook-test-state-prop-server",
+              name: room,
               host,
               protocol
             }}
             onAgent={(agent) => {
-              capturedAgent = agent;
+              observerAgent = agent;
+            }}
+          />
+          <TestAgentComponent
+            options={{
+              agent: "TestStateAgent",
+              name: room,
+              host,
+              protocol
+            }}
+            onAgent={(agent) => {
+              senderAgent = agent;
             }}
           />
         </SuspenseWrapper>
       );
 
       await vi.waitFor(
         () => {
-          expect(capturedAgent?.identified).toBe(true);
+          expect(observerAgent?.identified).toBe(true);
+          expect(senderAgent?.identified).toBe(true);
           const stateEl = container.querySelector(
             '[data-testid="agent-state"]'
           );
@@ -406,15 +420,16 @@ describe("useAgent hook", () => {
         { timeout: 10000 }
       );
 
-      // Send state — server will broadcast back, which updates agent.state
+      // Trigger a server-side setState. Unlike client setState, this exercises
+      // the server broadcast path back to the same connection.
       const newState = {
         count: 999,
         items: ["server-state"],
         lastUpdated: 2000
       };
-      capturedAgent!.setState(newState);
+      senderAgent!.setState(newState);
 
-      // Wait for the server broadcast to update state (second render)
+      // Wait for the server broadcast to update state (second render).
       await vi.waitFor(
         () => {
           const stateEl = container.querySelector(

packages/ai-chat/src/tests/message-concurrency.test.ts

@@ -380,10 +380,14 @@ describe("AIChatAgent messageConcurrency", () => {
       room
     );
 
+    // Keep the first turn in-flight long enough for both overlapping submits
+    // to reach the concurrency controller before the queued debounce turn can
+    // evaluate whether it was superseded. Otherwise slow CI WebSocket dispatch
+    // can let req-debounce-2 run before req-debounce-3 has been admitted.
     sendChatRequest(ws, "req-debounce-1", [firstUserMessage], {
       format: "plaintext",
-      chunkCount: 8,
-      chunkDelayMs: 80
+      chunkCount: 15,
+      chunkDelayMs: 150
     });
     await delay(50);
 
@@ -846,10 +850,13 @@ describe("AIChatAgent messageConcurrency", () => {
       room
     );
 
+    // Keep req-1 open long enough for both overlapping merge submits to be
+    // admitted before req-2 can acquire the turn lock and check supersession.
+    // This mirrors the primary merge strategy test above.
     sendChatRequest(ws, "req-resp-merge-1", [firstUserMessage], {
       format: "plaintext",
-      chunkCount: 8,
-      chunkDelayMs: 80
+      chunkCount: 15,
+      chunkDelayMs: 150
     });
     await delay(50);