test: de-flake the two assertions ejecting merge groups (#1787, #1952)

packages/mcp-gateway/src/server.test.ts

@@ -1417,12 +1417,25 @@ describe("mcp gateway lifecycle compatibility", () => {
   it("keeps aggregate initialize available when one upstream is unhealthy", async () => {
     const alpha = await createStrictMcpUpstream([{ name: "search", description: "Alpha search" }]);
     const hanging = await createHangingUpstream();
+    // `GatewayState.upstreamTimeoutMs` is ONE global budget, not per-upstream,
+    // so this deadline applies to the healthy `alpha` as much as to `stuck`.
+    // With `failureThreshold: 1`, a single `alpha` response slower than the
+    // budget opens its breaker for the rest of the test and `tools/list`
+    // returns `[]`. At the previous 150 ms that left a local in-process
+    // upstream almost no margin: on a CPU-starved runner it lost the race and
+    // this test failed `expected [] to deeply equal [ 'alpha__search' ]`,
+    // ejecting the merge group for #1952 on 2026-09-21.
+    //
+    // `hanging` never answers at all, so raising the budget does not weaken
+    // what this test checks — it only costs wall-clock. Reproduced directly:
+    // squeezing this to 1 ms reproduces that exact assertion locally.
+    const UPSTREAM_TIMEOUT_MS = 1_000;
     const gateway = await createAggregateGateway(
       {
         alpha: { url: alpha.url, credentialHeaders: [] },
         stuck: { url: hanging.url, credentialHeaders: [] },
       },
-      { timeoutMs: 150, failureThreshold: 1 },
+      { timeoutMs: UPSTREAM_TIMEOUT_MS, failureThreshold: 1 },
     );
 
     const start = Date.now();
@@ -1432,7 +1445,11 @@ describe("mcp gateway lifecycle compatibility", () => {
 
     expect(initialize.status).toBe(200);
     expect(clientSessionId).toBeTruthy();
-    expect(elapsed).toBeLessThan(1000);
+    // Proportional to the budget above, not an absolute wall-clock number:
+    // what this asserts is "initialize returns near the timeout rather than
+    // hanging on `stuck`", and an absolute bound re-introduces exactly the
+    // steal sensitivity this change removes.
+    expect(elapsed).toBeLessThan(UPSTREAM_TIMEOUT_MS * 5);
     expect(gateway.state.breaker.stateOf("stuck")).toBe("open");
 
     const list = await postJson(

ui/src/components/AgentActionButtons.test.tsx

@@ -124,17 +124,28 @@ describe("AgentActionButtons", () => {
     vi.clearAllMocks();
   });
 
-  function render(agent: Agent) {
-    root = createRoot(container);
-    root.render(
-      <QueryClientProvider client={queryClient}>
-        <AgentActionButtons agent={agent} companyId="company-1" runLabel="Run Heartbeat" />
-      </QueryClientProvider>,
-    );
+  // `root.render` must run INSIDE `act`, which is what every other UI suite
+  // here does (Inbox, Routines, CompanyInvites, ProjectDetail). Rendering
+  // outside it and relying on a later `flushReact()` to catch up makes the
+  // first paint depend on the scheduler getting a turn within that helper's
+  // single microtask + one `setTimeout(0)`. On a CPU-starved runner it does
+  // not, and the assertion reads `container.textContent === ""` — the
+  // `expected '' to contain 'Pause'` failure that ejected the merge group for
+  // #1787 on 2026-09-21. Inside `act`, the render is flushed before this
+  // returns, so the assertion no longer races the scheduler.
+  async function render(agent: Agent) {
+    await act(async () => {
+      root = createRoot(container);
+      root.render(
+        <QueryClientProvider client={queryClient}>
+          <AgentActionButtons agent={agent} companyId="company-1" runLabel="Run Heartbeat" />
+        </QueryClientProvider>,
+      );
+    });
   }
 
   it("replaces the pause slot with Clear error for error agents", async () => {
-    render(makeAgent({ status: "error" }));
+    await render(makeAgent({ status: "error" }));
     await flushReact();
 
     expect(container.textContent).toContain("Clear error");
@@ -150,7 +161,7 @@ describe("AgentActionButtons", () => {
   });
 
   it("calls clearError and refreshes agent-related queries", async () => {
-    render(makeAgent({ status: "error" }));
+    await render(makeAgent({ status: "error" }));
     await flushReact();
 
     await act(async () => {
@@ -169,7 +180,7 @@ describe("AgentActionButtons", () => {
   });
 
   it("keeps the normal pause action for non-error agents", async () => {
-    render(makeAgent({ status: "active" }));
+    await render(makeAgent({ status: "active" }));
     await flushReact();
 
     expect(container.textContent).toContain("Pause");