fix(testing): isolate bun module mocks across tests

src/features/background-agent/abort-with-timeout.test.ts

@@ -1,4 +1,4 @@
-import { afterAll, describe, expect, mock, test } from "bun:test"
+import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test"
 
 const logMock = mock(() => {})
 
@@ -20,6 +20,10 @@ function createClient(abort: (...args: Array<unknown>) => Promise<unknown>): Ope
 }
 
 describe("abortWithTimeout", () => {
+  beforeEach(() => {
+    logMock.mockClear()
+  })
+
   afterAll(() => {
     mock.restore()
   })

src/testing/module-mock-lifecycle-bun.test.ts

@@ -0,0 +1,161 @@
+/// <reference types="bun-types" />
+
+import { describe, expect, mock, test } from "bun:test"
+import { installModuleMockLifecycle } from "./module-mock-lifecycle"
+import { mockNestedFixture } from "./module-mock-lifecycle-nested/caller"
+
+describe("installModuleMockLifecycle Bun integration", () => {
+  test("restores in-test ESM mock bindings during cleanup", async () => {
+    // given
+    mock.module("./module-mock-lifecycle-fixture", () => ({ named: "mocked-aftereach" }))
+    const mockedModule = await import("./module-mock-lifecycle-fixture")
+
+    // when
+    mock.restore()
+
+    // then
+    expect(mockedModule.named).toBe("original")
+  })
+
+  test("keeps restored module mocks restored after test setup cleanup", async () => {
+    // when
+    const restoredModule = await import(`./module-mock-lifecycle-fixture?aftereach=${Date.now()}`)
+
+    // then
+    expect(restoredModule.named).toBe("original")
+  })
+
+  test("uses Bun resolution so extensionless TypeScript imports capture the loaded module", () => {
+    // given
+    const moduleCalls: Array<{ specifier: string; value: Record<string, unknown> }> = []
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        moduleCalls.push({ specifier, value: factory() })
+      },
+      restore: mock(() => {}),
+    }
+
+    installModuleMockLifecycle(mockApi)
+
+    // when
+    mockApi.module("./module-mock-lifecycle-fixture", () => ({ named: "mocked" }))
+    mockApi.restore()
+
+    // then
+    const restoreCall = moduleCalls.at(1)
+    expect(restoreCall?.specifier).toMatch(/\/src\/testing\/module-mock-lifecycle-fixture\.ts$/)
+    expect(restoreCall?.value.named).toBe("original")
+  })
+
+  test("resolves relative mocks from the helper caller instead of the lifecycle module", () => {
+    // given
+    const moduleCalls: Array<{ specifier: string; value: Record<string, unknown> }> = []
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        moduleCalls.push({ specifier, value: factory() })
+      },
+      restore: mock(() => {}),
+    }
+
+    installModuleMockLifecycle(mockApi)
+
+    // when
+    mockNestedFixture(mockApi)
+    mockApi.restore()
+
+    // then
+    const restoreCall = moduleCalls.at(1)
+    expect(restoreCall?.specifier).toMatch(/\/src\/testing\/module-mock-lifecycle-nested\/fixture\.ts$/)
+    expect(restoreCall?.value.named).toBe("original")
+  })
+
+  test("restores Bun module mocks for later cache-busted imports", async () => {
+    // given
+    mock.module("./module-mock-lifecycle-fixture", () => ({ named: "mocked" }))
+    const mockedModule = await import("./module-mock-lifecycle-fixture")
+
+    // when
+    mock.restore()
+    const restoredModule = await import(`./module-mock-lifecycle-fixture?restored=${Date.now()}`)
+
+    // then
+    expect(mockedModule.named).toBe("original")
+    expect(restoredModule.named).toBe("original")
+  })
+
+  test("keeps Bun module mocks restored when cleanup calls mock.restore twice", async () => {
+    // given
+    mock.module("./module-mock-lifecycle-fixture", () => ({ named: "mocked" }))
+    const mockedModule = await import("./module-mock-lifecycle-fixture")
+
+    // when
+    mock.restore()
+    mock.restore()
+    const restoredModule = await import(`./module-mock-lifecycle-fixture?double-restore=${Date.now()}`)
+
+    // then
+    expect(mockedModule.named).toBe("original")
+    expect(restoredModule.named).toBe("original")
+  })
+
+  test("allows a later relative mock after restoring a resolved original", async () => {
+    // given
+    mock.module("./module-mock-lifecycle-fixture", () => ({ named: "first mock" }))
+    const firstMockedModule = await import("./module-mock-lifecycle-fixture")
+
+    // when
+    mock.restore()
+    mock.module("./module-mock-lifecycle-fixture", () => ({ named: "second mock" }))
+    const secondMockedModule = await import("./module-mock-lifecycle-fixture")
+
+    // then
+    expect(firstMockedModule.named).toBe("second mock")
+    expect(secondMockedModule.named).toBe("second mock")
+  })
+
+  test("allows a later caller-relative mock after restoring a file-url original", async () => {
+    // given
+    mock.module("./module-mock-lifecycle-nested/fixture", () => ({ named: "first mock" }))
+    const firstMockedModule = await import("./module-mock-lifecycle-nested/fixture")
+
+    // when
+    mock.restore()
+    mockNestedFixture(mock)
+    const secondMockedModule = await import("./module-mock-lifecycle-nested/fixture")
+
+    // then
+    expect(firstMockedModule.named).toBe("mocked")
+    expect(secondMockedModule.named).toBe("mocked")
+  })
+
+  test("restores Bun module mocks outside the lifecycle helper directory", async () => {
+    // given
+    mock.module("../cli/run/json-output", () => ({
+      createJsonOutputManager: () => ({ mocked: true }),
+    }))
+    const mockedModule = await import("../cli/run/json-output")
+
+    // when
+    mock.restore()
+    const restoredModule = await import(`../cli/run/json-output?restored=${Date.now()}`)
+
+    // then
+    expect("mocked" in mockedModule.createJsonOutputManager()).toBe(false)
+    expect("mocked" in restoredModule.createJsonOutputManager()).toBe(false)
+    expect(typeof restoredModule.createJsonOutputManager().redirectToStderr).toBe("function")
+  })
+
+  test("restores dependency mocks after a cache-busted consumer import", async () => {
+    // given
+    mock.module("./module-mock-lifecycle-fixture", () => ({ named: "mocked" }))
+    const consumerModule = await import(`./module-mock-lifecycle-consumer?mocked=${Date.now()}`)
+
+    // when
+    mock.restore()
+    const restoredModule = await import(`./module-mock-lifecycle-fixture?restored-dep=${Date.now()}`)
+
+    // then
+    expect(consumerModule.consumed).toBe("original")
+    expect(restoredModule.named).toBe("original")
+  })
+})

src/testing/module-mock-lifecycle-consumer.ts

@@ -0,0 +1 @@
+export { named as consumed } from "./module-mock-lifecycle-fixture"

src/testing/module-mock-lifecycle-fixture.ts

@@ -0,0 +1 @@
+export const named = "original"

src/testing/module-mock-lifecycle-nested/caller.ts

@@ -0,0 +1,7 @@
+type MockModuleApi = {
+  module: (specifier: string, factory: () => Record<string, unknown>) => unknown
+}
+
+export function mockNestedFixture(mockApi: MockModuleApi): void {
+  mockApi.module("./fixture", () => ({ named: "mocked" }))
+}

src/testing/module-mock-lifecycle-nested/fixture.ts

@@ -0,0 +1 @@
+export const named = "original"

src/testing/module-mock-lifecycle-tracking.test.ts

@@ -0,0 +1,155 @@
+/// <reference types="bun-types" />
+
+import { describe, expect, mock, test } from "bun:test"
+import { installModuleMockLifecycle } from "./module-mock-lifecycle"
+
+describe("installModuleMockLifecycle active-test tracking", () => {
+  test("reapplies module-evaluation mocks with their mock factory after active restore", () => {
+    // given
+    const events: string[] = []
+    const loadOriginalModule = mock(() => ({ ok: true as const, value: { named: "original" } }))
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        events.push(`module:${specifier}:${String(factory().named)}`)
+      },
+      restore: mock(() => {
+        events.push("delegate:restore")
+      }),
+    }
+
+    const { beginTestMockTracking, endTestMockTracking } = installModuleMockLifecycle(mockApi, {
+      getCallerStack: () => "Error\n    at file:///repo/tests/top-level.test.ts:5:1\n    at moduleEvaluation (native:1:11)",
+      getCallerUrl: () => "file:///repo/tests/top-level.test.ts",
+      trackOnlyDuringActiveTest: true,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule,
+    })
+
+    // when
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+    beginTestMockTracking()
+    mockApi.restore()
+    endTestMockTracking()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:mocked",
+      "module:resolved:./dependency:mocked",
+    ])
+    expect(loadOriginalModule).not.toHaveBeenCalled()
+  })
+
+  test("restores original exports for active in-test mocks", () => {
+    // given
+    const events: string[] = []
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        events.push(`module:${specifier}:${String(factory().named)}`)
+      },
+      restore: mock(() => {
+        events.push("delegate:restore")
+      }),
+    }
+
+    const { beginTestMockTracking, endTestMockTracking } = installModuleMockLifecycle(mockApi, {
+      getCallerStack: () => "Error\n    at file:///repo/tests/example.test.ts:5:1\n    at test (native:1:11)",
+      getCallerUrl: () => "file:///repo/tests/example.test.ts",
+      trackOnlyDuringActiveTest: true,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    // when
+    beginTestMockTracking()
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+    mockApi.restore()
+    endTestMockTracking()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:original",
+      "module:resolved:./dependency:original",
+    ])
+  })
+
+  test("clears persistent module-evaluation snapshots when restore runs while inactive", () => {
+    // given
+    const events: string[] = []
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        events.push(`module:${specifier}:${String(factory().named)}`)
+      },
+      restore: mock(() => {
+        events.push("delegate:restore")
+      }),
+    }
+
+    const { beginTestMockTracking, endTestMockTracking } = installModuleMockLifecycle(mockApi, {
+      getCallerStack: () => "Error\n    at file:///repo/tests/top-level.test.ts:5:1\n    at moduleEvaluation (native:1:11)",
+      getCallerUrl: () => "file:///repo/tests/top-level.test.ts",
+      trackOnlyDuringActiveTest: true,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    // when
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+    mockApi.restore()
+    beginTestMockTracking()
+    mockApi.restore()
+    endTestMockTracking()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+      "delegate:restore",
+    ])
+  })
+
+  test("keeps unrelated module-evaluation mocks after inactive restore", () => {
+    // given
+    const events: string[] = []
+    let callerStack = "Error\n    at file:///repo/tests/first.test.ts:5:1\n    at moduleEvaluation (native:1:11)"
+    let callerUrl = "file:///repo/tests/first.test.ts"
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        events.push(`module:${specifier}:${String(factory().named)}`)
+      },
+      restore: mock(() => {
+        events.push("delegate:restore")
+      }),
+    }
+
+    installModuleMockLifecycle(mockApi, {
+      getCallerStack: () => callerStack,
+      getCallerUrl: () => callerUrl,
+      trackOnlyDuringActiveTest: true,
+      resolveSpecifier: (specifier, ownerUrl) => `resolved:${ownerUrl}:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    mockApi.module("./first", () => ({ named: "first top-level" }))
+    callerStack = "Error\n    at file:///repo/tests/second.test.ts:5:1\n    at moduleEvaluation (native:1:11)"
+    callerUrl = "file:///repo/tests/second.test.ts"
+    mockApi.module("./second", () => ({ named: "second top-level" }))
+
+    // when
+    callerStack = "Error\n    at file:///repo/tests/first.test.ts:10:1\n    at cleanup (native:1:11)"
+    callerUrl = "file:///repo/tests/first.test.ts"
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./first:first top-level",
+      "module:./second:second top-level",
+      "delegate:restore",
+      "module:./second:second top-level",
+      "module:resolved:file:///repo/tests/second.test.ts:./second:second top-level",
+    ])
+  })
+})

src/testing/module-mock-lifecycle.test.ts

@@ -4,9 +4,12 @@ import { describe, expect, mock, test } from "bun:test"
 import { installModuleMockLifecycle } from "./module-mock-lifecycle"
 
 describe("installModuleMockLifecycle", () => {
-  test("restores the original module export object instead of a cloned snapshot", () => {
+  test("preserves plain export identity while cloning ESM namespace snapshots", () => {
     // given
     const originalExports = { named: "original" }
+    const moduleNamespaceExports = Object.defineProperty({ named: "namespace" }, Symbol.toStringTag, {
+      value: "Module",
+    })
     const moduleCalls: Array<{ specifier: string; value: Record<string, unknown> }> = []
     const mockApi = {
       module: (specifier: string, factory: () => Record<string, unknown>) => {
@@ -18,17 +21,25 @@ describe("installModuleMockLifecycle", () => {
     installModuleMockLifecycle(mockApi, {
       getCallerUrl: () => "file:///repo/tests/example.test.ts",
       resolveSpecifier: (specifier) => `resolved:${specifier}`,
-      loadOriginalModule: () => ({ ok: true, value: originalExports }),
+      loadOriginalModule: (specifier) => ({
+        ok: true,
+        value: specifier === "./namespace" ? moduleNamespaceExports : originalExports,
+      }),
     })
 
     // when
     mockApi.module("./dependency", () => ({ named: "mocked" }))
+    mockApi.module("./namespace", () => ({ named: "mocked namespace" }))
     mockApi.restore()
 
     // then
-    expect(moduleCalls.map((call) => call.specifier)).toEqual(["./dependency", "resolved:./dependency"])
-    const restoreCall = moduleCalls.find((call) => call.specifier === "resolved:./dependency")
-    expect(restoreCall?.value).toBe(originalExports)
+    const plainRestoreCall = moduleCalls.find((call) => call.specifier === "./dependency" && call.value.named === "original")
+    const namespaceRestoreCall = moduleCalls.find(
+      (call) => call.specifier === "resolved:./namespace" && call.value.named === "namespace",
+    )
+    expect(plainRestoreCall?.value).toBe(originalExports)
+    expect(namespaceRestoreCall?.value).toEqual({ named: "namespace" })
+    expect(namespaceRestoreCall?.value).not.toBe(moduleNamespaceExports)
   })
 
   test("clears tracked snapshots after the delegate restore runs", () => {
@@ -58,6 +69,42 @@ describe("installModuleMockLifecycle", () => {
     expect(events).toEqual([
       "module:./dependency:mocked",
       "delegate:restore",
+      "module:./dependency:original",
+      "module:resolved:./dependency:original",
+    ])
+  })
+
+  test("reapplies the last restore snapshot when mock.restore is called twice", () => {
+    // given
+    const events: string[] = []
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        events.push(`module:${specifier}:${String(factory().named)}`)
+      },
+      restore: mock(() => {
+        events.push("delegate:restore")
+      }),
+    }
+
+    installModuleMockLifecycle(mockApi, {
+      getCallerUrl: () => "file:///repo/tests/example.test.ts",
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    // when
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+    mockApi.restore()
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:original",
+      "module:resolved:./dependency:original",
+      "delegate:restore",
+      "module:./dependency:original",
       "module:resolved:./dependency:original",
     ])
   })
@@ -87,6 +134,37 @@ describe("installModuleMockLifecycle", () => {
     expect(loadCount).toBe(1)
   })
 
+  test("restores each original mock specifier while sharing one resolved snapshot", () => {
+    // given
+    const moduleCalls: Array<{ specifier: string; value: Record<string, unknown> }> = []
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        moduleCalls.push({ specifier, value: factory() })
+      },
+      restore: mock(() => {}),
+    }
+
+    installModuleMockLifecycle(mockApi, {
+      getCallerUrl: () => "file:///repo/tests/example.test.ts",
+      resolveSpecifier: () => "file:///repo/src/dependency.ts",
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    // when
+    mockApi.module("../src/dependency", () => ({ named: "first mock" }))
+    mockApi.module("./dependency", () => ({ named: "second mock" }))
+    mockApi.restore()
+
+    // then
+    expect(moduleCalls).toEqual([
+      { specifier: "../src/dependency", value: { named: "first mock" } },
+      { specifier: "./dependency", value: { named: "second mock" } },
+      { specifier: "../src/dependency", value: { named: "original" } },
+      { specifier: "file:///repo/src/dependency.ts", value: { named: "original" } },
+      { specifier: "./dependency", value: { named: "original" } },
+    ])
+  })
+
   test("does not restore unresolved modules to avoid cleanup errors", () => {
     // given
     const moduleCalls: Array<{ specifier: string; value: Record<string, unknown> }> = []
@@ -107,7 +185,7 @@ describe("installModuleMockLifecycle", () => {
     mockApi.module("virtual:missing", () => ({ named: "mocked" }))
     mockApi.restore()
 
-    // then - only the original mock call, no restore call for unresolved module
+    // then
     expect(moduleCalls).toEqual([{ specifier: "virtual:missing", value: { named: "mocked" } }])
   })
 })

src/testing/module-mock-lifecycle.ts

@@ -1,5 +1,6 @@
 import { createRequire } from "node:module"
-import { pathToFileURL } from "node:url"
+import { fileURLToPath } from "node:url"
+import { defaultGetCallerStack, isModuleEvaluationStack, resolveCallerUrlFromStack } from "./module-mock-stack"
 
 type MockModuleFactory = () => Record<string, unknown>
 
@@ -13,28 +14,43 @@ type ModuleLoadResult =
   | { ok: false; error: Error }
 
 type ModuleSnapshot = {
-  restoreSpecifier: string
+  restoreOriginalSpecifiers: boolean
+  restoreSpecifiers: Set<string>
+  restoreFactory: MockModuleFactory
+}
+
+type PersistentModuleSnapshot = {
+  ownerUrls: Set<string>
+  restoreSpecifiers: Set<string>
   restoreFactory: MockModuleFactory
 }
 
 type ModuleMockLifecycleOptions = {
+  getCallerStack?: () => string
   getCallerUrl?: () => string
+  trackOnlyDuringActiveTest?: boolean
   resolveSpecifier?: (specifier: string, callerUrl: string) => string
   loadOriginalModule?: (specifier: string, callerUrl: string) => ModuleLoadResult
 }
 
-function toError(error: unknown): Error {
-  if (error instanceof Error) {
-    return error
-  }
+let originalLoadNonce = 0
 
-  return new Error(String(error))
+function toError(error: unknown): Error {
+  return error instanceof Error ? error : new Error(String(error))
 }
 
 function isModuleExports(moduleValue: unknown): moduleValue is Record<string, unknown> {
   return moduleValue !== null && typeof moduleValue === "object"
 }
 
+function isModuleNamespaceObject(moduleValue: Record<string, unknown>): boolean {
+  return Object.prototype.toString.call(moduleValue) === "[object Module]"
+}
+
+function shouldRestoreOriginalSpecifier(moduleValue: unknown): boolean {
+  return !isModuleExports(moduleValue) || !isModuleNamespaceObject(moduleValue)
+}
+
 function createRestoreExports(moduleValue: unknown): Record<string, unknown> {
   if (typeof moduleValue === "function") {
     const functionExports = Object.assign({}, moduleValue)
@@ -45,47 +61,42 @@ function createRestoreExports(moduleValue: unknown): Record<string, unknown> {
   }
 
   if (isModuleExports(moduleValue)) {
+    if (isModuleNamespaceObject(moduleValue)) {
+      return { ...moduleValue }
+    }
+
     return moduleValue
   }
 
   return { default: moduleValue }
 }
 
-function normalizeStackPath(rawPath: string): string {
-  if (rawPath.startsWith("file://")) {
-    return rawPath
-  }
-
-  return pathToFileURL(rawPath).href
+function resolveWithBun(specifier: string, callerUrl: string): string {
+  const callerDirectory = fileURLToPath(new URL(".", callerUrl))
+  return Bun.resolveSync(specifier, callerDirectory)
 }
 
-function defaultGetCallerUrl(): string {
-  const stack = new Error().stack ?? ""
-  const lines = stack.split("\n")
-
-  for (const line of lines) {
-    const match = line.match(/(?:\()?(file:\/\/[^\s)]+|\/[^\s):]+):(\d+):(\d+)/)
-    const candidatePath = match?.[1]
-    if (!candidatePath) {
-      continue
-    }
-
-    if (
-      candidatePath.includes("/test-setup.ts") ||
-      candidatePath.includes("/src/testing/module-mock-lifecycle.ts")
-    ) {
-      continue
-    }
-
-    return normalizeStackPath(candidatePath)
-  }
-
-  return import.meta.url
+function isSchemeSpecifier(specifier: string): boolean {
+  return /^[a-zA-Z][a-zA-Z\d+.-]*:/.test(specifier)
 }
 
 function defaultResolveSpecifier(specifier: string, callerUrl: string): string {
   try {
-    return import.meta.resolve(specifier, callerUrl)
+    return resolveWithBun(specifier, callerUrl)
+  } catch {
+    return specifier
+  }
+}
+
+function createOriginalLoadSpecifier(specifier: string, callerUrl: string): string {
+  try {
+    const resolved = resolveWithBun(specifier, callerUrl)
+    if (isSchemeSpecifier(resolved)) {
+      return specifier
+    }
+
+    originalLoadNonce += 1
+    return `${resolved}?omo_original=${originalLoadNonce}`
   } catch {
     return specifier
   }
@@ -94,7 +105,7 @@ function defaultResolveSpecifier(specifier: string, callerUrl: string): string {
 function defaultLoadOriginalModule(specifier: string, callerUrl: string): ModuleLoadResult {
   try {
     const require = createRequire(callerUrl)
-    return { ok: true, value: require(specifier) }
+    return { ok: true, value: require(createOriginalLoadSpecifier(specifier, callerUrl)) }
   } catch (error) {
     return { ok: false, error: toError(error) }
   }
@@ -103,46 +114,141 @@ function defaultLoadOriginalModule(specifier: string, callerUrl: string): Module
 export function installModuleMockLifecycle(
   mockApi: MockApi,
   options: ModuleMockLifecycleOptions = {},
-): { restoreModuleMocks: () => void } {
+): {
+  beginTestMockTracking: () => void
+  endTestMockTracking: () => void
+  restoreModuleMocks: () => void
+} {
   const snapshots = new Map<string, ModuleSnapshot>()
+  const persistentSnapshots = new Map<string, PersistentModuleSnapshot>()
+  let lastRestoredSnapshots: ModuleSnapshot[] = []
+  let isActiveTest = !options.trackOnlyDuringActiveTest
   const delegateModule = mockApi.module.bind(mockApi)
   const delegateRestore = mockApi.restore.bind(mockApi)
-  const getCallerUrl = options.getCallerUrl ?? defaultGetCallerUrl
+  const getCallerStack = options.getCallerStack ?? defaultGetCallerStack
   const resolveSpecifier = options.resolveSpecifier ?? defaultResolveSpecifier
   const loadOriginalModule = options.loadOriginalModule ?? defaultLoadOriginalModule
 
-  function restoreModuleMocks(): void {
-    for (const snapshot of snapshots.values()) {
-      delegateModule(snapshot.restoreSpecifier, snapshot.restoreFactory)
+  function getCallerUrl(callerStack: string): string {
+    return options.getCallerUrl?.() ?? resolveCallerUrlFromStack(callerStack)
+  }
+
+  function restoreModuleMocksForRestoreCall(): void {
+    const snapshotsToRestore = snapshots.size > 0 ? Array.from(snapshots.values()) : lastRestoredSnapshots
+
+    for (const snapshot of snapshotsToRestore) {
+      for (const restoreSpecifier of snapshot.restoreSpecifiers) {
+        delegateModule(restoreSpecifier, snapshot.restoreFactory)
+      }
     }
 
-    snapshots.clear()
+    if (snapshots.size > 0) {
+      lastRestoredSnapshots = snapshotsToRestore
+      snapshots.clear()
+    }
+  }
+
+  function restorePersistentModuleMocksForRestoreCall(): void {
+    for (const snapshot of persistentSnapshots.values()) {
+      for (const restoreSpecifier of snapshot.restoreSpecifiers) {
+        delegateModule(restoreSpecifier, snapshot.restoreFactory)
+      }
+    }
+  }
+
+  function clearPersistentModuleMocksForOwner(ownerUrl: string): void {
+    for (const [resolvedSpecifier, snapshot] of persistentSnapshots) {
+      snapshot.ownerUrls.delete(ownerUrl)
+      if (snapshot.ownerUrls.size === 0) {
+        persistentSnapshots.delete(resolvedSpecifier)
+      }
+    }
+  }
+
+  function restoreModuleMocks(): void {
+    if (snapshots.size === 0) {
+      return
+    }
+
+    restoreModuleMocksForRestoreCall()
+  }
+
+  function beginTestMockTracking(): void {
+    isActiveTest = true
+  }
+
+  function endTestMockTracking(): void {
+    isActiveTest = !options.trackOnlyDuringActiveTest
   }
 
   mockApi.module = (specifier: string, factory: MockModuleFactory): unknown => {
-    const callerUrl = getCallerUrl()
-    const restoreSpecifier = resolveSpecifier(specifier, callerUrl)
+    lastRestoredSnapshots = []
+    const callerStack = getCallerStack()
+    const callerUrl = getCallerUrl(callerStack)
+    const isModuleEvaluation = isModuleEvaluationStack(callerStack)
 
-    if (!snapshots.has(restoreSpecifier)) {
-      const originalModule = loadOriginalModule(specifier, callerUrl)
-
-      if (originalModule.ok) {
-        const restoreExports = createRestoreExports(originalModule.value)
-        snapshots.set(restoreSpecifier, {
-          restoreSpecifier,
-          restoreFactory: () => restoreExports,
+    if (isModuleEvaluation) {
+      const resolvedSpecifier = resolveSpecifier(specifier, callerUrl)
+      const existingSnapshot = persistentSnapshots.get(resolvedSpecifier)
+      if (existingSnapshot) {
+        existingSnapshot.ownerUrls.add(callerUrl)
+        existingSnapshot.restoreSpecifiers.add(specifier)
+        existingSnapshot.restoreSpecifiers.add(resolvedSpecifier)
+      } else {
+        persistentSnapshots.set(resolvedSpecifier, {
+          ownerUrls: new Set([callerUrl]),
+          restoreSpecifiers: new Set([specifier, resolvedSpecifier]),
+          restoreFactory: factory,
         })
       }
+      return delegateModule(specifier, factory)
+    }
+
+    if (isActiveTest) {
+      const resolvedSpecifier = resolveSpecifier(specifier, callerUrl)
+      const existingSnapshot = snapshots.get(resolvedSpecifier)
+
+      if (existingSnapshot) {
+        if (existingSnapshot.restoreOriginalSpecifiers) {
+          existingSnapshot.restoreSpecifiers.add(specifier)
+        }
+        existingSnapshot.restoreSpecifiers.add(resolvedSpecifier)
+      } else {
+        const originalModule = loadOriginalModule(specifier, callerUrl)
+
+        if (originalModule.ok) {
+          const restoreExports = createRestoreExports(originalModule.value)
+          const restoreOriginalSpecifiers = shouldRestoreOriginalSpecifier(originalModule.value)
+          snapshots.set(resolvedSpecifier, {
+            restoreOriginalSpecifiers,
+            restoreSpecifiers: new Set(
+              restoreOriginalSpecifiers ? [specifier, resolvedSpecifier] : [resolvedSpecifier],
+            ),
+            restoreFactory: () => restoreExports,
+          })
+        }
+      }
     }
 
     return delegateModule(specifier, factory)
   }
 
   mockApi.restore = (): unknown => {
+    const callerStack = getCallerStack()
+    const callerUrl = getCallerUrl(callerStack)
     const result = delegateRestore()
-    restoreModuleMocks()
+    if (!isActiveTest) {
+      snapshots.clear()
+      lastRestoredSnapshots = []
+      clearPersistentModuleMocksForOwner(callerUrl)
+      restorePersistentModuleMocksForRestoreCall()
+      return result
+    }
+
+    restoreModuleMocksForRestoreCall()
+    restorePersistentModuleMocksForRestoreCall()
     return result
   }
 
-  return { restoreModuleMocks }
+  return { beginTestMockTracking, endTestMockTracking, restoreModuleMocks }
 }

src/testing/module-mock-stack.ts

@@ -0,0 +1,61 @@
+import { pathToFileURL } from "node:url"
+
+function normalizeStackPath(rawPath: string): string {
+  return rawPath.startsWith("file://") ? rawPath : pathToFileURL(rawPath).href
+}
+
+function getStackCandidatePath(line: string): string | undefined {
+  const match = line.match(/(?:\()?(file:\/\/[^\s)]+|\/[^\s):]+):(\d+):(\d+)/)
+  return match?.[1]
+}
+
+function isIgnoredCallerPath(candidatePath: string): boolean {
+  return (
+    candidatePath.includes("/test-setup.ts") ||
+    candidatePath.includes("/src/testing/module-mock-lifecycle.ts") ||
+    candidatePath.includes("/src/testing/module-mock-stack.ts")
+  )
+}
+
+export function defaultGetCallerStack(): string {
+  return new Error().stack ?? ""
+}
+
+export function resolveCallerUrlFromStack(stack: string): string {
+  const lines = stack.split("\n")
+
+  for (const line of lines) {
+    const candidatePath = getStackCandidatePath(line)
+    if (!candidatePath) {
+      continue
+    }
+
+    if (isIgnoredCallerPath(candidatePath)) {
+      continue
+    }
+
+    return normalizeStackPath(candidatePath)
+  }
+
+  return import.meta.url
+}
+
+export function isModuleEvaluationStack(stack: string): boolean {
+  const lines = stack.split("\n")
+
+  for (const [index, line] of lines.entries()) {
+    const candidatePath = getStackCandidatePath(line)
+    if (!candidatePath || isIgnoredCallerPath(candidatePath)) {
+      continue
+    }
+
+    const nextFrame = lines[index + 1] ?? ""
+    return (
+      nextFrame.includes("moduleEvaluation") ||
+      nextFrame.includes("asyncModuleEvaluation") ||
+      nextFrame.includes("loadAndEvaluateModule")
+    )
+  }
+
+  return false
+}

test-setup.ts

@@ -10,7 +10,9 @@ import { getOmoOpenCodeCacheDir } from "./src/shared/data-path"
 import { releaseAllPromptAsyncReservationsForTesting } from "./src/shared/prompt-async-gate"
 import { installModuleMockLifecycle } from "./src/testing/module-mock-lifecycle"
 
-const { restoreModuleMocks } = installModuleMockLifecycle(mock)
+const { beginTestMockTracking, endTestMockTracking, restoreModuleMocks } = installModuleMockLifecycle(mock, {
+  trackOnlyDuringActiveTest: true,
+})
 let environmentSnapshot: NodeJS.ProcessEnv = { ...process.env }
 let workingDirectorySnapshot = process.cwd()
 
@@ -23,6 +25,7 @@ function cleanupRulesInjectorStorage(): void {
 }
 
 beforeEach(() => {
+  beginTestMockTracking()
   environmentSnapshot = { ...process.env }
   workingDirectorySnapshot = process.cwd()
   process.env.OMO_DISABLE_POSTHOG = "true"
@@ -65,4 +68,5 @@ afterEach(() => {
   releaseAllPromptAsyncReservationsForTesting()
   mock.restore()
   restoreModuleMocks()
+  endTestMockTracking()
 })