Iterate on hooks implementation

features/enterprise/api/src/main/kotlin/io/element/android/features/enterprise/api/ClientBuilderEnterpriseHook.kt

@@ -13,6 +13,23 @@ import io.element.android.libraries.matrix.api.core.SessionId
 /**
  * A hook that can be used to customize the [MatrixClientBuilder] for enterprise features.
  */
-fun interface ClientBuilderEnterpriseHook {
-    suspend operator fun invoke(clientBuilder: MatrixClientBuilder, sessionId: SessionId): MatrixClientBuilder
+interface ClientBuilderEnterpriseHook {
+    /**
+     * Customize the [MatrixClientBuilder] for enterprise features.
+     * This method is invoked everytime a new [MatrixClientBuilder] is created.
+     *
+     * @param clientBuilder The [MatrixClientBuilder] to customize.
+     * @return The customized [MatrixClientBuilder].
+     */
+    suspend fun tweakClientBuilder(clientBuilder: MatrixClientBuilder): MatrixClientBuilder
+
+    /**
+     * Customize the [MatrixClientBuilder] for enterprise features.
+     * This method is invoked when a new [MatrixClientBuilder] is created to build a client for a specific session.
+     *
+     * @param clientBuilder The [MatrixClientBuilder] to customize.
+     * @param sessionId The [SessionId] for which the [MatrixClientBuilder] is being created.
+     * @return The customized [MatrixClientBuilder].
+     */
+    suspend fun tweakClientBuilder(clientBuilder: MatrixClientBuilder, sessionId: SessionId): MatrixClientBuilder
 }

features/enterprise/impl-foss/src/main/kotlin/io/element/android/features/enterprise/impl/DefaultClientBuilderEnterpriseHook.kt

@@ -18,7 +18,12 @@ import io.element.android.libraries.matrix.api.core.SessionId
  */
 @ContributesBinding(AppScope::class)
 class DefaultClientBuilderEnterpriseHook : ClientBuilderEnterpriseHook {
-    override suspend fun invoke(clientBuilder: MatrixClientBuilder, sessionId: SessionId): MatrixClientBuilder {
+    override suspend fun tweakClientBuilder(clientBuilder: MatrixClientBuilder): MatrixClientBuilder {
+        // No modification
+        return clientBuilder
+    }
+
+    override suspend fun tweakClientBuilder(clientBuilder: MatrixClientBuilder, sessionId: SessionId): MatrixClientBuilder {
         // No modification
         return clientBuilder
     }

features/enterprise/test/src/main/kotlin/io/element/android/features/enterprise/test/FakeClientBuilderEnterpriseHook.kt

@@ -0,0 +1,25 @@
+/*
+ * Copyright (c) 2026 Element Creations Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial.
+ * Please see LICENSE files in the repository root for full details.
+ */
+
+package io.element.android.features.enterprise.test
+
+import io.element.android.features.enterprise.api.ClientBuilderEnterpriseHook
+import io.element.android.libraries.matrix.api.MatrixClientBuilder
+import io.element.android.libraries.matrix.api.core.SessionId
+
+class FakeClientBuilderEnterpriseHook(
+    private val tweakClientBuilderResult: (MatrixClientBuilder) -> MatrixClientBuilder = { it },
+    private val tweakClientBuilderWithSessionIdResult: (MatrixClientBuilder, SessionId) -> MatrixClientBuilder = { clientBuilder, _ -> clientBuilder },
+) : ClientBuilderEnterpriseHook {
+    override suspend fun tweakClientBuilder(clientBuilder: MatrixClientBuilder): MatrixClientBuilder {
+        return tweakClientBuilderResult(clientBuilder)
+    }
+
+    override suspend fun tweakClientBuilder(clientBuilder: MatrixClientBuilder, sessionId: SessionId): MatrixClientBuilder {
+        return tweakClientBuilderWithSessionIdResult(clientBuilder, sessionId)
+    }
+}

libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/RustMatrixClientFactory.kt

@@ -99,7 +99,9 @@ class RustMatrixClientFactory(
         )
             .homeserverUrl(sessionData.homeserverUrl)
             .enableAutomaticBackPagination(featureFlagService.isFeatureEnabled(FeatureFlags.AutomaticBackPagination))
-            .let { (clientBuilderEnterpriseHook(RustMatrixClientBuilder(it), SessionId(sessionData.userId)) as RustMatrixClientBuilder).inner }
+            .let {
+                (clientBuilderEnterpriseHook.tweakClientBuilder(RustMatrixClientBuilder(it), SessionId(sessionData.userId)) as RustMatrixClientBuilder).inner
+            }
             .use { it.build() }
 
         client.setMediaRetentionPolicy(
@@ -219,6 +221,9 @@ class RustMatrixClientFactory(
                 // Workaround for non-nullable proxy parameter in the SDK, since each call to the ClientBuilder returns a new reference we need to keep
                 proxyProvider.provides()?.let { proxy(it) } ?: this
             }
+            .let {
+                (clientBuilderEnterpriseHook.tweakClientBuilder(RustMatrixClientBuilder(it)) as RustMatrixClientBuilder).inner
+            }
     }
 }
 

libraries/matrix/impl/src/test/kotlin/io/element/android/libraries/matrix/impl/RustMatrixClientFactoryTest.kt

@@ -9,6 +9,7 @@
 package io.element.android.libraries.matrix.impl
 
 import com.google.common.truth.Truth.assertThat
+import io.element.android.features.enterprise.test.FakeClientBuilderEnterpriseHook
 import io.element.android.libraries.featureflag.test.FakeFeatureFlagService
 import io.element.android.libraries.matrix.api.core.SessionId
 import io.element.android.libraries.matrix.impl.auth.FakeProxyProvider
@@ -65,5 +66,5 @@ fun TestScope.createRustMatrixClientFactory(
     clientBuilderProvider = clientBuilderProvider,
     sqliteStoreBuilderProvider = FakeSqliteStoreBuilderProvider(),
     workManagerScheduler = workManagerScheduler,
-    clientBuilderEnterpriseHook = { builder, _ -> builder },
+    clientBuilderEnterpriseHook = FakeClientBuilderEnterpriseHook(),
 )

libraries/matrix/impl/src/test/kotlin/io/element/android/libraries/matrix/impl/RustTemporaryMatrixClientFactoryTest.kt

@@ -8,6 +8,7 @@
 package io.element.android.libraries.matrix.impl
 
 import com.google.common.truth.Truth.assertThat
+import io.element.android.features.enterprise.test.FakeClientBuilderEnterpriseHook
 import io.element.android.libraries.featureflag.test.FakeFeatureFlagService
 import io.element.android.libraries.matrix.impl.auth.FakeProxyProvider
 import io.element.android.libraries.matrix.impl.paths.SessionPathsFactory
@@ -77,6 +78,6 @@ class RustTemporaryMatrixClientFactoryTest {
         clientBuilderProvider = clientBuilderProvider,
         sqliteStoreBuilderProvider = FakeSqliteStoreBuilderProvider(),
         workManagerScheduler = workManagerScheduler,
-        clientBuilderEnterpriseHook = { builder, _ -> builder },
+        clientBuilderEnterpriseHook = FakeClientBuilderEnterpriseHook(),
     )
 }