diff --git a/src/auto-reply/reply/conversation-label-generator.test.ts b/src/auto-reply/reply/conversation-label-generator.test.ts
index fe77e179dab..674e0feb422 100644
--- a/src/auto-reply/reply/conversation-label-generator.test.ts
+++ b/src/auto-reply/reply/conversation-label-generator.test.ts
@@ -7,8 +7,9 @@ const resolveDefaultModelForAgent = vi.hoisted(() => vi.fn());
 const resolveModelAsync = vi.hoisted(() => vi.fn());
 const prepareModelForSimpleCompletion = vi.hoisted(() => vi.fn());
 
-vi.mock("@mariozechner/pi-ai", async (importOriginal) => {
-  const original = await importOriginal<typeof import("@mariozechner/pi-ai")>();
+vi.mock("@mariozechner/pi-ai", async () => {
+  const original =
+    await vi.importActual<typeof import("@mariozechner/pi-ai")>("@mariozechner/pi-ai");
   return {
     ...original,
     completeSimple,