fix(testing): restore mock specifier aliases

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

@@ -42,7 +42,7 @@ describe("installModuleMockLifecycle Bun integration", () => {
     mockApi.restore()
 
     // then
-    const restoreCall = moduleCalls.at(1)
+    const restoreCall = moduleCalls.find((call) => /\/src\/testing\/module-mock-lifecycle-fixture\.ts$/.test(call.specifier))
     expect(restoreCall?.specifier).toMatch(/\/src\/testing\/module-mock-lifecycle-fixture\.ts$/)
     expect(restoreCall?.value.named).toBe("original")
   })
@@ -64,7 +64,7 @@ describe("installModuleMockLifecycle Bun integration", () => {
     mockApi.restore()
 
     // then
-    const restoreCall = moduleCalls.at(1)
+    const restoreCall = moduleCalls.find((call) => /\/src\/testing\/module-mock-lifecycle-nested\/fixture\.ts$/.test(call.specifier))
     expect(restoreCall?.specifier).toMatch(/\/src\/testing\/module-mock-lifecycle-nested\/fixture\.ts$/)
     expect(restoreCall?.value.named).toBe("original")
   })

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

@@ -21,6 +21,7 @@ describe("installModuleMockLifecycle active-test tracking", () => {
       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,
+      isPersistentModuleMockOwner: () => true,
       resolveSpecifier: (specifier) => `resolved:${specifier}`,
       loadOriginalModule,
     })
@@ -76,6 +77,42 @@ describe("installModuleMockLifecycle active-test tracking", () => {
     ])
   })
 
+  test("restores active in-test mocks even when Bun stack includes module evaluation", () => {
+    // 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 moduleEvaluation (native:1:11)",
+      getCallerUrl: () => "file:///repo/tests/example.test.ts",
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => 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[] = []
@@ -92,6 +129,7 @@ describe("installModuleMockLifecycle active-test tracking", () => {
       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,
+      isPersistentModuleMockOwner: () => true,
       resolveSpecifier: (specifier) => `resolved:${specifier}`,
       loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
     })
@@ -111,6 +149,164 @@ describe("installModuleMockLifecycle active-test tracking", () => {
     ])
   })
 
+  test("does not restore ordinary inactive mocks before their owner test cleanup runs", () => {
+    // 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/late-loaded.test.ts:5:1\n    at moduleEvaluation (native:1:11)",
+      getCallerUrl: () => "file:///repo/tests/late-loaded.test.ts",
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => true,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    beginTestMockTracking()
+    endTestMockTracking()
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+
+    // when
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+    ])
+  })
+
+  test("restores original exports for beforeAll-style inactive mocks after owner test cleanup", () => {
+    // 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/before-all.test.ts:5:1\n    at beforeAll (native:1:11)",
+      getCallerUrl: () => "file:///repo/tests/before-all.test.ts",
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => true,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+
+    // when
+    beginTestMockTracking()
+    mockApi.restore()
+    endTestMockTracking()
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:mocked",
+      "module:resolved:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:original",
+      "module:resolved:./dependency:original",
+    ])
+  })
+
+  test("restores original exports for persistent module-evaluation mocks after owner test cleanup", () => {
+    // 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,
+      isPersistentModuleMockOwner: () => true,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule,
+    })
+
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+
+    // when
+    beginTestMockTracking()
+    mockApi.restore()
+    endTestMockTracking()
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:mocked",
+      "module:resolved:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:original",
+      "module:resolved:./dependency:original",
+    ])
+    expect(loadOriginalModule).toHaveBeenCalledTimes(1)
+  })
+
+  test("reapplies the last active restore snapshot when a later inactive restore runs", () => {
+    // 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 cleanup (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()
+    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",
+    ])
+  })
+
   test("keeps unrelated module-evaluation mocks after inactive restore", () => {
     // given
     const events: string[] = []
@@ -129,6 +325,7 @@ describe("installModuleMockLifecycle active-test tracking", () => {
       getCallerStack: () => callerStack,
       getCallerUrl: () => callerUrl,
       trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => true,
       resolveSpecifier: (specifier, ownerUrl) => `resolved:${ownerUrl}:${specifier}`,
       loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
     })
@@ -152,4 +349,128 @@ describe("installModuleMockLifecycle active-test tracking", () => {
       "module:resolved:file:///repo/tests/second.test.ts:./second:second top-level",
     ])
   })
+
+  test("keeps module-evaluation mocks when inactive restore cannot resolve the owner", () => {
+    // given
+    const events: string[] = []
+    let callerStack = "Error\n    at file:///repo/tests/top-level.test.ts:5:1\n    at moduleEvaluation (native:1:11)"
+    let callerUrl = "file:///repo/tests/top-level.test.ts"
+    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: () => callerStack,
+      getCallerUrl: () => callerUrl,
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => true,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+
+    // when
+    callerStack = "Error\n    at nativeAfterAll (native:1:11)"
+    callerUrl = "file:///repo/testing/unknown-owner.ts"
+    mockApi.restore()
+    beginTestMockTracking()
+    mockApi.restore()
+    endTestMockTracking()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:mocked",
+      "module:resolved:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:mocked",
+      "module:resolved:./dependency:mocked",
+    ])
+  })
+
+  test("restores re-applied persistent mocks when inactive restore cannot resolve the owner", () => {
+    // given
+    const events: string[] = []
+    let callerStack = "Error\n    at file:///repo/tests/top-level.test.ts:5:1\n    at moduleEvaluation (native:1:11)"
+    let callerUrl = "file:///repo/tests/top-level.test.ts"
+    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: () => callerStack,
+      getCallerUrl: () => callerUrl,
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => true,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+    beginTestMockTracking()
+    mockApi.restore()
+    endTestMockTracking()
+
+    // when
+    callerStack = "Error\n    at nativeAfterAll (native:1:11)"
+    callerUrl = "file:///repo/testing/unknown-owner.ts"
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:mocked",
+      "module:resolved:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:original",
+      "module:resolved:./dependency:original",
+    ])
+  })
+
+  test("does not reapply ordinary module-evaluation mocks after active restore", () => {
+    // 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/file-local.test.ts:5:1\n    at moduleEvaluation (native:1:11)",
+      getCallerUrl: () => "file:///repo/tests/file-local.test.ts",
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => false,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    // when
+    mockApi.module("./dependency", () => ({ named: "file local" }))
+    beginTestMockTracking()
+    mockApi.restore()
+    endTestMockTracking()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:file local",
+      "delegate:restore",
+    ])
+  })
 })

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

@@ -4,7 +4,7 @@ import { describe, expect, mock, test } from "bun:test"
 import { installModuleMockLifecycle } from "./module-mock-lifecycle"
 
 describe("installModuleMockLifecycle", () => {
-  test("preserves plain export identity while cloning ESM namespace snapshots", () => {
+  test("restores every mock specifier while cloning ESM namespace snapshots", () => {
     // given
     const originalExports = { named: "original" }
     const moduleNamespaceExports = Object.defineProperty({ named: "namespace" }, Symbol.toStringTag, {
@@ -34,12 +34,17 @@ describe("installModuleMockLifecycle", () => {
 
     // then
     const plainRestoreCall = moduleCalls.find((call) => call.specifier === "./dependency" && call.value.named === "original")
-    const namespaceRestoreCall = moduleCalls.find(
+    const namespaceOriginalRestoreCall = moduleCalls.find(
+      (call) => call.specifier === "./namespace" && call.value.named === "namespace",
+    )
+    const namespaceResolvedRestoreCall = moduleCalls.find(
       (call) => call.specifier === "resolved:./namespace" && call.value.named === "namespace",
     )
     expect(plainRestoreCall?.value).toBe(originalExports)
-    expect(namespaceRestoreCall?.value).toEqual({ named: "namespace" })
+    expect(namespaceOriginalRestoreCall?.value).toEqual({ named: "namespace" })
-    expect(namespaceRestoreCall?.value).not.toBe(moduleNamespaceExports)
+    expect(namespaceOriginalRestoreCall?.value).not.toBe(moduleNamespaceExports)
+    expect(namespaceResolvedRestoreCall?.value).toEqual({ named: "namespace" })
+    expect(namespaceResolvedRestoreCall?.value).not.toBe(moduleNamespaceExports)
   })
 
   test("clears tracked snapshots after the delegate restore runs", () => {

src/testing/module-mock-lifecycle.ts

@@ -1,6 +1,6 @@
 import { createRequire } from "node:module"
 import { fileURLToPath } from "node:url"
-import { defaultGetCallerStack, isModuleEvaluationStack, resolveCallerUrlFromStack } from "./module-mock-stack"
+import { defaultGetCallerStack, resolveCallerUrlFromStack } from "./module-mock-stack"
 
 type MockModuleFactory = () => Record<string, unknown>
 
@@ -14,13 +14,13 @@ type ModuleLoadResult =
   | { ok: false; error: Error }
 
 type ModuleSnapshot = {
-  restoreOriginalSpecifiers: boolean
   restoreSpecifiers: Set<string>
   restoreFactory: MockModuleFactory
 }
 
 type PersistentModuleSnapshot = {
-  ownerUrls: Set<string>
+  originalSpecifier: string
+  reappliedDuringActiveRestore: boolean
   restoreSpecifiers: Set<string>
   restoreFactory: MockModuleFactory
 }
@@ -29,6 +29,7 @@ type ModuleMockLifecycleOptions = {
   getCallerStack?: () => string
   getCallerUrl?: () => string
   trackOnlyDuringActiveTest?: boolean
+  isPersistentModuleMockOwner?: (callerUrl: string) => boolean
   resolveSpecifier?: (specifier: string, callerUrl: string) => string
   loadOriginalModule?: (specifier: string, callerUrl: string) => ModuleLoadResult
 }
@@ -47,10 +48,6 @@ 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)
@@ -111,6 +108,10 @@ function defaultLoadOriginalModule(specifier: string, callerUrl: string): Module
   }
 }
 
+function defaultIsPersistentModuleMockOwner(_callerUrl: string): boolean {
+  return true
+}
+
 export function installModuleMockLifecycle(
   mockApi: MockApi,
   options: ModuleMockLifecycleOptions = {},
@@ -120,14 +121,16 @@ export function installModuleMockLifecycle(
   restoreModuleMocks: () => void
 } {
   const snapshots = new Map<string, ModuleSnapshot>()
-  const persistentSnapshots = new Map<string, PersistentModuleSnapshot>()
+  const persistentSnapshots = new Map<string, Map<string, PersistentModuleSnapshot>>()
   let lastRestoredSnapshots: ModuleSnapshot[] = []
   let isActiveTest = !options.trackOnlyDuringActiveTest
+  let hasStartedTest = false
   const delegateModule = mockApi.module.bind(mockApi)
   const delegateRestore = mockApi.restore.bind(mockApi)
   const getCallerStack = options.getCallerStack ?? defaultGetCallerStack
   const resolveSpecifier = options.resolveSpecifier ?? defaultResolveSpecifier
   const loadOriginalModule = options.loadOriginalModule ?? defaultLoadOriginalModule
+  const isPersistentModuleMockOwner = options.isPersistentModuleMockOwner ?? defaultIsPersistentModuleMockOwner
 
   function getCallerUrl(callerStack: string): string {
     return options.getCallerUrl?.() ?? resolveCallerUrlFromStack(callerStack)
@@ -148,18 +151,66 @@ export function installModuleMockLifecycle(
     }
   }
 
-  function restorePersistentModuleMocksForRestoreCall(): void {
+  function restorePersistentModuleMocksForRestoreCall(markReapplied: boolean): void {
-    for (const snapshot of persistentSnapshots.values()) {
+    for (const snapshotsByOwner of persistentSnapshots.values()) {
-      for (const restoreSpecifier of snapshot.restoreSpecifiers) {
+      for (const snapshot of snapshotsByOwner.values()) {
-        delegateModule(restoreSpecifier, snapshot.restoreFactory)
+        if (markReapplied) {
+          snapshot.reappliedDuringActiveRestore = true
+        }
+        for (const restoreSpecifier of snapshot.restoreSpecifiers) {
+          delegateModule(restoreSpecifier, snapshot.restoreFactory)
+        }
       }
     }
   }
 
-  function clearPersistentModuleMocksForOwner(ownerUrl: string): void {
+  function restorePersistentOriginals(snapshot: PersistentModuleSnapshot, ownerUrl: string): void {
-    for (const [resolvedSpecifier, snapshot] of persistentSnapshots) {
+    const originalModule = loadOriginalModule(snapshot.originalSpecifier, ownerUrl)
-      snapshot.ownerUrls.delete(ownerUrl)
+    if (!originalModule.ok) {
-      if (snapshot.ownerUrls.size === 0) {
+      return
+    }
+
+    const originalFactory = () => createRestoreExports(originalModule.value)
+    for (const restoreSpecifier of snapshot.restoreSpecifiers) {
+      delegateModule(restoreSpecifier, originalFactory)
+    }
+  }
+
+  function clearPersistentModuleMocksForOwner(
+    ownerUrl: string,
+    restoreOriginals: boolean,
+    forceRestoreOriginals = false,
+  ): void {
+    let clearedOwnerSnapshot = false
+
+    for (const [resolvedSpecifier, snapshotsByOwner] of persistentSnapshots) {
+      const snapshot = snapshotsByOwner.get(ownerUrl)
+      if (snapshot) {
+        if (restoreOriginals && (forceRestoreOriginals || snapshot.reappliedDuringActiveRestore)) {
+          restorePersistentOriginals(snapshot, ownerUrl)
+        }
+        snapshotsByOwner.delete(ownerUrl)
+        clearedOwnerSnapshot = true
+      }
+      if (snapshotsByOwner.size === 0) {
+        persistentSnapshots.delete(resolvedSpecifier)
+      }
+    }
+
+    if (clearedOwnerSnapshot || !restoreOriginals) {
+      return
+    }
+
+    for (const [resolvedSpecifier, snapshotsByOwner] of persistentSnapshots) {
+      for (const [snapshotOwnerUrl, snapshot] of snapshotsByOwner) {
+        if (!snapshot.reappliedDuringActiveRestore) {
+          continue
+        }
+
+        restorePersistentOriginals(snapshot, snapshotOwnerUrl)
+        snapshotsByOwner.delete(snapshotOwnerUrl)
+      }
+      if (snapshotsByOwner.size === 0) {
         persistentSnapshots.delete(resolvedSpecifier)
       }
     }
@@ -174,6 +225,7 @@ export function installModuleMockLifecycle(
   }
 
   function beginTestMockTracking(): void {
+    hasStartedTest = true
     isActiveTest = true
   }
 
@@ -185,22 +237,24 @@ export function installModuleMockLifecycle(
     lastRestoredSnapshots = []
     const callerStack = getCallerStack()
     const callerUrl = getCallerUrl(callerStack)
-    const isModuleEvaluation = isModuleEvaluationStack(callerStack)
 
-    if (isModuleEvaluation) {
+    if (!isActiveTest && isPersistentModuleMockOwner(callerUrl)) {
       const resolvedSpecifier = resolveSpecifier(specifier, callerUrl)
-      const existingSnapshot = persistentSnapshots.get(resolvedSpecifier)
+      const snapshotsByOwner = persistentSnapshots.get(resolvedSpecifier) ?? new Map<string, PersistentModuleSnapshot>()
+      const existingSnapshot = snapshotsByOwner.get(callerUrl)
       if (existingSnapshot) {
-        existingSnapshot.ownerUrls.add(callerUrl)
         existingSnapshot.restoreSpecifiers.add(specifier)
         existingSnapshot.restoreSpecifiers.add(resolvedSpecifier)
+        existingSnapshot.restoreFactory = factory
       } else {
-        persistentSnapshots.set(resolvedSpecifier, {
+        snapshotsByOwner.set(callerUrl, {
-          ownerUrls: new Set([callerUrl]),
+          originalSpecifier: specifier,
+          reappliedDuringActiveRestore: false,
           restoreSpecifiers: new Set([specifier, resolvedSpecifier]),
           restoreFactory: factory,
         })
       }
+      persistentSnapshots.set(resolvedSpecifier, snapshotsByOwner)
       return delegateModule(specifier, factory)
     }
 
@@ -209,21 +263,15 @@ export function installModuleMockLifecycle(
       const existingSnapshot = snapshots.get(resolvedSpecifier)
 
       if (existingSnapshot) {
-        if (existingSnapshot.restoreOriginalSpecifiers) {
+        existingSnapshot.restoreSpecifiers.add(specifier)
-          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([specifier, resolvedSpecifier]),
-            restoreSpecifiers: new Set(
-              restoreOriginalSpecifiers ? [specifier, resolvedSpecifier] : [resolvedSpecifier],
-            ),
             restoreFactory: () => restoreExports,
           })
         }
@@ -238,15 +286,16 @@ export function installModuleMockLifecycle(
     const callerUrl = getCallerUrl(callerStack)
     const result = delegateRestore()
     if (!isActiveTest) {
+      restoreModuleMocksForRestoreCall()
       snapshots.clear()
       lastRestoredSnapshots = []
-      clearPersistentModuleMocksForOwner(callerUrl)
+      clearPersistentModuleMocksForOwner(callerUrl, hasStartedTest)
-      restorePersistentModuleMocksForRestoreCall()
+      restorePersistentModuleMocksForRestoreCall(false)
       return result
     }
 
     restoreModuleMocksForRestoreCall()
-    restorePersistentModuleMocksForRestoreCall()
+    restorePersistentModuleMocksForRestoreCall(true)
     return result
   }