diff --git a/ui/src/components/DocumentAnnotationLayer.test.tsx b/ui/src/components/DocumentAnnotationLayer.test.tsx
index cec5c12317b..955a51a8e5f 100644
--- a/ui/src/components/DocumentAnnotationLayer.test.tsx
+++ b/ui/src/components/DocumentAnnotationLayer.test.tsx
@@ -1,6 +1,5 @@
 // @vitest-environment jsdom
 
-import { act } from "react";
 import { createRoot, type Root } from "react-dom/client";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { DocumentAnnotationLayer } from "./DocumentAnnotationLayer";
@@ -16,6 +15,12 @@ vi.mock("@/lib/document-annotation-selection", () => ({
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 (globalThis as any).IS_REACT_ACT_ENVIRONMENT = true;
 
+async function act(callback: () => void | Promise<void>) {
+  await callback();
+  await Promise.resolve();
+  await new Promise((resolve) => setTimeout(resolve, 0));
+}
+
 function makeRect(left: number, top: number, width: number, height: number): DOMRect {
   return {
     x: left,
@@ -49,9 +54,9 @@ describe("DocumentAnnotationLayer", () => {
     rectSpy = vi.spyOn(HTMLElement.prototype, "getBoundingClientRect").mockReturnValue(makeRect(0, 0, 400, 300));
   });
 
-  afterEach(() => {
+  afterEach(async () => {
     if (root) {
-      act(() => root?.unmount());
+      await act(() => root?.unmount());
       root = null;
     }
     rectSpy.mockRestore();
diff --git a/ui/src/components/IssueDocumentAnnotations.test.tsx b/ui/src/components/IssueDocumentAnnotations.test.tsx
index 5e94deecde1..8206956547c 100644
--- a/ui/src/components/IssueDocumentAnnotations.test.tsx
+++ b/ui/src/components/IssueDocumentAnnotations.test.tsx
@@ -1,6 +1,6 @@
 // @vitest-environment jsdom
 
-import { act, useState } from "react";
+import { useState } from "react";
 import { createRoot } from "react-dom/client";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import type {
@@ -94,10 +94,14 @@ vi.mock("./DocumentAnnotationLayer", () => ({
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 (globalThis as any).IS_REACT_ACT_ENVIRONMENT = true;
 
+async function act(callback: () => void | Promise<void>) {
+  await callback();
+  await Promise.resolve();
+  await new Promise((resolve) => setTimeout(resolve, 0));
+}
+
 async function flush() {
-  await act(async () => {
-    await new Promise((resolve) => setTimeout(resolve, 0));
-  });
+  await act(() => {});
 }
 
 function setTextareaValue(textarea: HTMLTextAreaElement, value: string) {