diff --git a/platform/lang-impl/src/com/intellij/openapi/module/impl/AutomaticModuleUnloaderImpl.kt b/platform/lang-impl/src/com/intellij/openapi/module/impl/AutomaticModuleUnloaderImpl.kt
index c69932055bd7..96b9f58c3a86 100644
--- a/platform/lang-impl/src/com/intellij/openapi/module/impl/AutomaticModuleUnloaderImpl.kt
+++ b/platform/lang-impl/src/com/intellij/openapi/module/impl/AutomaticModuleUnloaderImpl.kt
@@ -19,6 +19,8 @@ import com.intellij.platform.workspace.jps.entities.ModuleEntity
 import com.intellij.platform.workspace.jps.entities.ModuleId
 import com.intellij.platform.workspace.storage.EntityStorage
 import com.intellij.platform.workspace.storage.MutableEntityStorage
+import com.intellij.platform.workspace.storage.instrumentation.EntityStorageInstrumentationApi
+import com.intellij.platform.workspace.storage.instrumentation.MutableEntityStorageInstrumentation
 import com.intellij.util.xmlb.annotations.XCollection
 import com.intellij.workspaceModel.ide.impl.legacyBridge.module.getModuleLevelLibraries
 import com.intellij.xml.util.XmlStringUtil
@@ -31,6 +33,7 @@ import kotlinx.coroutines.launch
 @State(name = "AutomaticModuleUnloader", storages = [(Storage(StoragePathMacros.WORKSPACE_FILE))])
 internal class AutomaticModuleUnloaderImpl(private val project: Project) : SimplePersistentStateComponent<LoadedModulesListStorage>(LoadedModulesListStorage()),
                                                                            AutomaticModuleUnloader {
+  @OptIn(EntityStorageInstrumentationApi::class)
   override fun calculateNewModules(currentModules: Set<String>, builder: MutableEntityStorage, unloadedEntityBuilder: MutableEntityStorage): Pair<List<String>, List<String>> {
     val defaultResult = Pair(emptyList<String>(), emptyList<String>())
 
@@ -78,7 +81,7 @@ internal class AutomaticModuleUnloaderImpl(private val project: Project) : Simpl
       val moduleEntitiesToAdd = snapshot.entities(ModuleEntity::class.java).filter { it.name in toUnloadSet }.toList()
       val moduleEntitiesToRemove = builder.entities(ModuleEntity::class.java).filter { it.name in toUnloadSet }.toList()
       for (moduleEntity in moduleEntitiesToAdd) {
-        unloadedEntityBuilder.addEntity(moduleEntity)
+        (unloadedEntityBuilder as MutableEntityStorageInstrumentation).addEntity(moduleEntity)
         moduleEntity.getModuleLevelLibraries(snapshot).forEach { libraryEntity ->
           unloadedEntityBuilder.addEntity(libraryEntity)
         }
diff --git a/platform/projectModel-impl/src/com/intellij/workspaceModel/ide/impl/legacyBridge/module/ModuleManagerBridgeImpl.kt b/platform/projectModel-impl/src/com/intellij/workspaceModel/ide/impl/legacyBridge/module/ModuleManagerBridgeImpl.kt
index 6b8d59327570..2aba2c5bcc10 100644
--- a/platform/projectModel-impl/src/com/intellij/workspaceModel/ide/impl/legacyBridge/module/ModuleManagerBridgeImpl.kt
+++ b/platform/projectModel-impl/src/com/intellij/workspaceModel/ide/impl/legacyBridge/module/ModuleManagerBridgeImpl.kt
@@ -35,6 +35,8 @@ import com.intellij.platform.workspace.jps.JpsProjectFileEntitySource
 import com.intellij.platform.workspace.jps.entities.*
 import com.intellij.platform.workspace.jps.serialization.impl.ModulePath
 import com.intellij.platform.workspace.storage.*
+import com.intellij.platform.workspace.storage.instrumentation.EntityStorageInstrumentationApi
+import com.intellij.platform.workspace.storage.instrumentation.MutableEntityStorageInstrumentation
 import com.intellij.platform.workspace.storage.query.entities
 import com.intellij.platform.workspace.storage.query.map
 import com.intellij.platform.workspace.storage.url.VirtualFileUrl
@@ -399,6 +401,7 @@ abstract class ModuleManagerBridgeImpl(private val project: Project,
     }
   }
 
+  @OptIn(EntityStorageInstrumentationApi::class)
   private fun addAndRemoveModules(builder: MutableEntityStorage,
                                   entitiesToAdd: List<ModuleEntity>,
                                   entitiesToRemove: List<ModuleEntity>,
@@ -407,7 +410,7 @@ abstract class ModuleManagerBridgeImpl(private val project: Project,
       builder.removeEntity(entity)
     }
     for (entity in entitiesToAdd) {
-      builder.addEntity(entity)
+      (builder as MutableEntityStorageInstrumentation).addEntity(entity)
       entity.getModuleLevelLibraries(storageContainingEntitiesToAdd).forEach { libraryEntity ->
         builder.addEntity(libraryEntity)
       }
diff --git a/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/ExternalModuleImlFileEntitiesSerializer.kt b/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/ExternalModuleImlFileEntitiesSerializer.kt
index 3a9946009d8b..c4a96704f642 100644
--- a/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/ExternalModuleImlFileEntitiesSerializer.kt
+++ b/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/ExternalModuleImlFileEntitiesSerializer.kt
@@ -1,4 +1,4 @@
-// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.platform.workspace.jps.serialization.impl
 
 import com.intellij.openapi.diagnostic.logger
@@ -35,7 +35,7 @@ internal class ExternalModuleImlFileEntitiesSerializer(modulePath: ModulePath,
 
   override fun loadEntities(reader: JpsFileContentReader,
                             errorReporter: ErrorReporter,
-                            virtualFileManager: VirtualFileUrlManager): LoadingResult<Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity>>> {
+                            virtualFileManager: VirtualFileUrlManager): LoadingResult<Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity.Builder<out WorkspaceEntity>>>> {
     return LoadingResult(emptyMap(), null)
   }
 
diff --git a/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/JpsArtifactEntitiesSerializer.kt b/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/JpsArtifactEntitiesSerializer.kt
index 82d93cb8cd9e..8688d303454f 100644
--- a/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/JpsArtifactEntitiesSerializer.kt
+++ b/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/JpsArtifactEntitiesSerializer.kt
@@ -154,7 +154,7 @@ internal open class JpsArtifactEntitiesSerializer(override val fileUrl: VirtualF
     reader: JpsFileContentReader,
     errorReporter: ErrorReporter,
     virtualFileManager: VirtualFileUrlManager
-  ): LoadingResult<Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity>>> = loadEntitiesTimeMs.addMeasuredTime {
+  ): LoadingResult<Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity.Builder<out WorkspaceEntity>>>> = loadEntitiesTimeMs.addMeasuredTime {
     val artifactListElement = runCatchingXmlIssues { reader.loadComponent(fileUrl.url, ARTIFACT_MANAGER_COMPONENT_NAME) }
       .onFailure { return@addMeasuredTime LoadingResult(emptyMap(), it) }
       .getOrThrow()
@@ -200,7 +200,7 @@ internal open class JpsArtifactEntitiesSerializer(override val fileUrl: VirtualF
 
   override fun checkAndAddToBuilder(builder: MutableEntityStorage,
                                     orphanage: MutableEntityStorage,
-                                    newEntities: Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity>>) {
+                                    newEntities: Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity.Builder<out WorkspaceEntity>>>) {
     if (preserveOrder) {
       val order = newEntities[ArtifactsOrderEntity::class.java]?.singleOrNull() as? ArtifactsOrderEntity.Builder
       if (order != null) {
diff --git a/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/JpsLibraryEntitiesSerializer.kt b/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/JpsLibraryEntitiesSerializer.kt
index 7b2bafa5c369..595d844af88a 100644
--- a/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/JpsLibraryEntitiesSerializer.kt
+++ b/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/JpsLibraryEntitiesSerializer.kt
@@ -135,7 +135,7 @@ open class JpsLibraryEntitiesSerializer(override val fileUrl: VirtualFileUrl,
     reader: JpsFileContentReader,
     errorReporter: ErrorReporter,
     virtualFileManager: VirtualFileUrlManager
-  ): LoadingResult<Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity>>> = loadEntitiesTimeMs.addMeasuredTime {
+  ): LoadingResult<Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity.Builder<out WorkspaceEntity>>>> = loadEntitiesTimeMs.addMeasuredTime {
     val libraryTableTag = runCatchingXmlIssues { reader.loadComponent(fileUrl.url, LIBRARY_TABLE_COMPONENT_NAME) }
                             .onFailure { return@addMeasuredTime LoadingResult(emptyMap(), null) }
                             .getOrThrow() ?: return@addMeasuredTime LoadingResult(emptyMap(), null)
@@ -159,7 +159,7 @@ open class JpsLibraryEntitiesSerializer(override val fileUrl: VirtualFileUrl,
   @Suppress("UNCHECKED_CAST")
   override fun checkAndAddToBuilder(builder: MutableEntityStorage,
                                     orphanage: MutableEntityStorage,
-                                    newEntities: Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity>>) {
+                                    newEntities: Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity.Builder<out WorkspaceEntity>>>) {
     val libraries = (newEntities[LibraryEntity::class.java] as? List<LibraryEntity.Builder>) ?: emptyList()
     libraries.forEach {
       val symbolicId = LibraryId(it.name, it.tableId)
diff --git a/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/JpsSdkEntitySerializer.kt b/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/JpsSdkEntitySerializer.kt
index 0c74ccce2b5f..22d6f39db3c6 100644
--- a/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/JpsSdkEntitySerializer.kt
+++ b/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/JpsSdkEntitySerializer.kt
@@ -1,4 +1,4 @@
-// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.platform.workspace.jps.serialization.impl
 
 import com.intellij.openapi.diagnostic.logger
@@ -65,8 +65,11 @@ class JpsSdkEntitySerializer(val entitySource: JpsGlobalFileEntitySource, privat
     get() = SdkEntity::class.java
 
 
-  override fun loadEntities(reader: JpsFileContentReader, errorReporter: ErrorReporter,
-                            virtualFileManager: VirtualFileUrlManager): LoadingResult<Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity>>> {
+  override fun loadEntities(
+    reader: JpsFileContentReader,
+    errorReporter: ErrorReporter,
+    virtualFileManager: VirtualFileUrlManager,
+  ): LoadingResult<Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity.Builder<out WorkspaceEntity>>>> {
     val sdkTag = reader.loadComponent(entitySource.file.url, SDK_TABLE_COMPONENT_NAME) ?: return LoadingResult(emptyMap(), null)
     val sdkEntities = sdkTag.getChildren(ELEMENT_JDK).map { sdkElement -> loadSdkEntity(sdkElement, virtualFileManager) }
     return LoadingResult(mapOf(SdkEntity::class.java to sdkEntities))
@@ -121,8 +124,9 @@ class JpsSdkEntitySerializer(val entitySource: JpsGlobalFileEntitySource, privat
     return result
   }
 
-  override fun checkAndAddToBuilder(builder: MutableEntityStorage, orphanage: MutableEntityStorage,
-                                    newEntities: Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity>>) {
+  override fun checkAndAddToBuilder(builder: MutableEntityStorage,
+                                    orphanage: MutableEntityStorage,
+                                    newEntities: Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity.Builder<out WorkspaceEntity>>>) {
     newEntities.values.flatten().forEach { builder addEntity  it }
   }
 
diff --git a/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/ModuleImlFileEntitiesSerializer.kt b/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/ModuleImlFileEntitiesSerializer.kt
index 528ae8bbc8f0..cda5fb21f8aa 100644
--- a/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/ModuleImlFileEntitiesSerializer.kt
+++ b/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/ModuleImlFileEntitiesSerializer.kt
@@ -67,7 +67,7 @@ internal open class ModuleImlFileEntitiesSerializer(internal val modulePath: Mod
     reader: JpsFileContentReader,
     errorReporter: ErrorReporter,
     virtualFileManager: VirtualFileUrlManager
-  ): LoadingResult<Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity>>> = loadEntitiesTimeMs.addMeasuredTime {
+  ): LoadingResult<Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity.Builder<out WorkspaceEntity>>>> = loadEntitiesTimeMs.addMeasuredTime {
 
     val moduleLibrariesCollector: MutableMap<LibraryId, LibraryEntity.Builder> = HashMap()
     val newModuleEntity: ModuleEntity.Builder?
@@ -174,7 +174,7 @@ internal open class ModuleImlFileEntitiesSerializer(internal val modulePath: Mod
 
   override fun checkAndAddToBuilder(builder: MutableEntityStorage,
                                     orphanage: MutableEntityStorage,
-                                    newEntities: Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity>>) {
+                                    newEntities: Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity.Builder<out WorkspaceEntity>>>) {
 
     val (orphans, elements) = newEntities.values.asSequence().flatten().partition { it.entitySource is OrphanageWorkerEntitySource }
 
diff --git a/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/jpsFormatEntitiesSerialization.kt b/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/jpsFormatEntitiesSerialization.kt
index a26ef9fc0d1b..16e39db19d36 100644
--- a/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/jpsFormatEntitiesSerialization.kt
+++ b/platform/workspace/jps/src/com/intellij/platform/workspace/jps/serialization/impl/jpsFormatEntitiesSerialization.kt
@@ -53,8 +53,8 @@ interface JpsFileEntitiesSerializer<E : WorkspaceEntity> {
    */
   fun loadEntities(reader: JpsFileContentReader,
                    errorReporter: ErrorReporter,
-                   virtualFileManager: VirtualFileUrlManager): LoadingResult<Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity>>>
-  fun checkAndAddToBuilder(builder: MutableEntityStorage, orphanage: MutableEntityStorage, newEntities: Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity>>)
+                   virtualFileManager: VirtualFileUrlManager): LoadingResult<Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity.Builder<out WorkspaceEntity>>>>
+  fun checkAndAddToBuilder(builder: MutableEntityStorage, orphanage: MutableEntityStorage, newEntities: Map<Class<out WorkspaceEntity>, Collection<WorkspaceEntity.Builder<out WorkspaceEntity>>>)
 
   fun saveEntities(mainEntities: Collection<E>,
                    entities: Map<Class<out WorkspaceEntity>, List<WorkspaceEntity>>,
diff --git a/platform/workspace/storage/src/com/intellij/platform/workspace/storage/MutableEntityStorage.kt b/platform/workspace/storage/src/com/intellij/platform/workspace/storage/MutableEntityStorage.kt
index 7942607aa51d..726976e3ab40 100644
--- a/platform/workspace/storage/src/com/intellij/platform/workspace/storage/MutableEntityStorage.kt
+++ b/platform/workspace/storage/src/com/intellij/platform/workspace/storage/MutableEntityStorage.kt
@@ -138,7 +138,6 @@ public interface MutableEntityStorage : EntityStorage {
    * If any of the children exists in `this` storage, the reference to this child will be added instead of creating a new child entity.
    */
   public infix fun <M: WorkspaceEntity.Builder<T>, T : WorkspaceEntity> addEntity(entity: M): T
-  public infix fun <T : WorkspaceEntity> addEntity(entity: T): T
 
   /**
    * Modifies the given entity [e] by passing a builder interface for it to [change]. 
diff --git a/platform/workspace/storage/src/com/intellij/platform/workspace/storage/instrumentation/EntityStorageInstrumentation.kt b/platform/workspace/storage/src/com/intellij/platform/workspace/storage/instrumentation/EntityStorageInstrumentation.kt
index 9ce6eecf0b3d..db8e2126ce7f 100644
--- a/platform/workspace/storage/src/com/intellij/platform/workspace/storage/instrumentation/EntityStorageInstrumentation.kt
+++ b/platform/workspace/storage/src/com/intellij/platform/workspace/storage/instrumentation/EntityStorageInstrumentation.kt
@@ -128,6 +128,8 @@ public interface MutableEntityStorageInstrumentation : MutableEntityStorage, Ent
   public fun getManyChildrenBuilders(connectionId: ConnectionId, parent: WorkspaceEntity.Builder<*>): Sequence<WorkspaceEntity.Builder<*>>
 
   public fun getParentBuilder(connectionId: ConnectionId, child: WorkspaceEntity.Builder<*>): WorkspaceEntity.Builder<*>?
+
+  public infix fun <T : WorkspaceEntity> addEntity(entity: T): T
 }
 
 /**