Support loading SDKs for LocalController#loadSdk.

Introduce SdkRegistry concept - object responsible for lifecycle of particular SDKs.
Extract local SDKs loading logic from SdkSandboxManagerCompat to LocalSdkRegistry,
share it between SdkSandboxManagerCompat and LocalController.

Bug: 323825555
Test: LocalSdkRegistryTest
Test: LocalControllerTest, SdkSandboxManagerCompatTest
Change-Id: If0ad5989adc73ba83f69b5908d638832874e8d74
diff --git a/privacysandbox/sdkruntime/sdkruntime-client/src/androidTest/java/androidx/privacysandbox/sdkruntime/client/SdkSandboxManagerCompatTest.kt b/privacysandbox/sdkruntime/sdkruntime-client/src/androidTest/java/androidx/privacysandbox/sdkruntime/client/SdkSandboxManagerCompatTest.kt
index 81d35cf..af1cd4f 100644
--- a/privacysandbox/sdkruntime/sdkruntime-client/src/androidTest/java/androidx/privacysandbox/sdkruntime/client/SdkSandboxManagerCompatTest.kt
+++ b/privacysandbox/sdkruntime/sdkruntime-client/src/androidTest/java/androidx/privacysandbox/sdkruntime/client/SdkSandboxManagerCompatTest.kt
@@ -19,13 +19,10 @@
 import android.content.ContextWrapper
 import android.os.Binder
 import android.os.Bundle
-import androidx.privacysandbox.sdkruntime.client.activity.LocalSdkActivityHandlerRegistry
 import androidx.privacysandbox.sdkruntime.client.activity.SdkActivity
 import androidx.privacysandbox.sdkruntime.client.loader.CatchingSdkActivityHandler
 import androidx.privacysandbox.sdkruntime.client.loader.asTestSdk
-import androidx.privacysandbox.sdkruntime.client.loader.extractSdkProviderFieldValue
 import androidx.privacysandbox.sdkruntime.core.LoadSdkCompatException
-import androidx.privacysandbox.sdkruntime.core.LoadSdkCompatException.Companion.LOAD_SDK_INTERNAL_ERROR
 import androidx.privacysandbox.sdkruntime.core.LoadSdkCompatException.Companion.LOAD_SDK_SDK_DEFINED_ERROR
 import androidx.privacysandbox.sdkruntime.core.SandboxedSdkInfo
 import androidx.test.core.app.ActivityScenario
@@ -133,47 +130,6 @@
     }
 
     @Test
-    fun loadSdk_whenLocalSdkFailedToLoad_throwsInternalErrorException() {
-        val context = ApplicationProvider.getApplicationContext<Context>()
-        val managerCompat = SdkSandboxManagerCompat.from(context)
-
-        val result = assertThrows(LoadSdkCompatException::class.java) {
-            runBlocking {
-                managerCompat.loadSdk(
-                    TestSdkConfigs.forSdkName("invalidEntryPoint").packageName,
-                    Bundle()
-                )
-            }
-        }
-
-        assertThat(result.loadSdkErrorCode).isEqualTo(LOAD_SDK_INTERNAL_ERROR)
-        assertThat(result.message).isEqualTo("Failed to instantiate local SDK")
-    }
-
-    @Test
-    fun loadSdk_afterUnloading_loadSdkAgain() {
-        val context = ApplicationProvider.getApplicationContext<Context>()
-        val managerCompat = SdkSandboxManagerCompat.from(context)
-
-        val sdkName = TestSdkConfigs.CURRENT.packageName
-
-        val sdkToUnload = runBlocking {
-            managerCompat.loadSdk(sdkName, Bundle())
-        }
-
-        managerCompat.unloadSdk(sdkName)
-
-        val reloadedSdk = runBlocking {
-            managerCompat.loadSdk(sdkName, Bundle())
-        }
-
-        assertThat(managerCompat.getSandboxedSdks())
-            .containsExactly(reloadedSdk)
-        assertThat(reloadedSdk.getInterface())
-            .isNotEqualTo(sdkToUnload.getInterface())
-    }
-
-    @Test
     @SdkSuppress(maxSdkVersion = 33)
     fun unloadSdk_whenNoLocalSdkLoadedAndApiBelow34_doesntThrow() {
         val context = ApplicationProvider.getApplicationContext<Context>()
@@ -191,47 +147,13 @@
         runBlocking {
             managerCompat.loadSdk(sdkName, Bundle())
         }
-        val sdkProvider = managerCompat.getLocallyLoadedSdk(sdkName)!!.sdkProvider
-
         managerCompat.unloadSdk(sdkName)
 
-        val isBeforeUnloadSdkCalled = sdkProvider.extractSdkProviderFieldValue<Boolean>(
-            fieldName = "isBeforeUnloadSdkCalled"
-        )
-
-        assertThat(isBeforeUnloadSdkCalled)
-            .isTrue()
-
         assertThat(managerCompat.getSandboxedSdks())
             .isEmpty()
     }
 
     @Test
-    fun unloadSdk_unregisterActivityHandlers() {
-        val context = ApplicationProvider.getApplicationContext<Context>()
-        val managerCompat = SdkSandboxManagerCompat.from(context)
-
-        val packageName = TestSdkConfigs.forSdkName("v4").packageName
-        val localSdk = runBlocking {
-            managerCompat.loadSdk(
-                packageName,
-                Bundle()
-            )
-        }
-
-        val testSdk = localSdk.asTestSdk()
-        val token = testSdk.registerSdkSandboxActivityHandler(CatchingSdkActivityHandler())
-
-        val registeredBefore = LocalSdkActivityHandlerRegistry.isRegistered(token)
-        assertThat(registeredBefore).isTrue()
-
-        managerCompat.unloadSdk(packageName)
-
-        val registeredAfter = LocalSdkActivityHandlerRegistry.isRegistered(token)
-        assertThat(registeredAfter).isFalse()
-    }
-
-    @Test
     @SdkSuppress(maxSdkVersion = 33)
     fun addSdkSandboxProcessDeathCallback_whenApiBelow34_doesntThrow() {
         val context = ApplicationProvider.getApplicationContext<Context>()
diff --git a/privacysandbox/sdkruntime/sdkruntime-client/src/androidTest/java/androidx/privacysandbox/sdkruntime/client/controller/LocalControllerTest.kt b/privacysandbox/sdkruntime/sdkruntime-client/src/androidTest/java/androidx/privacysandbox/sdkruntime/client/controller/LocalControllerTest.kt
index 1d4260a..2c99ab4 100644
--- a/privacysandbox/sdkruntime/sdkruntime-client/src/androidTest/java/androidx/privacysandbox/sdkruntime/client/controller/LocalControllerTest.kt
+++ b/privacysandbox/sdkruntime/sdkruntime-client/src/androidTest/java/androidx/privacysandbox/sdkruntime/client/controller/LocalControllerTest.kt
@@ -19,11 +19,12 @@
 import android.os.Binder
 import android.os.Bundle
 import androidx.privacysandbox.sdkruntime.client.activity.LocalSdkActivityHandlerRegistry
-import androidx.privacysandbox.sdkruntime.client.loader.LocalSdkProvider
 import androidx.privacysandbox.sdkruntime.core.AppOwnedSdkSandboxInterfaceCompat
+import androidx.privacysandbox.sdkruntime.core.LoadSdkCompatException
 import androidx.privacysandbox.sdkruntime.core.SandboxedSdkCompat
 import androidx.privacysandbox.sdkruntime.core.activity.ActivityHolder
 import androidx.privacysandbox.sdkruntime.core.activity.SdkSandboxActivityHandlerCompat
+import androidx.privacysandbox.sdkruntime.core.controller.SdkSandboxControllerCompat
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.google.common.truth.Truth.assertThat
@@ -35,26 +36,51 @@
 @RunWith(AndroidJUnit4::class)
 class LocalControllerTest {
 
-    private lateinit var locallyLoadedSdks: LocallyLoadedSdks
+    private lateinit var localSdkRegistry: StubLocalSdkRegistry
     private lateinit var appOwnedSdkRegistry: StubAppOwnedSdkInterfaceRegistry
     private lateinit var controller: LocalController
 
     @Before
     fun setUp() {
-        locallyLoadedSdks = LocallyLoadedSdks()
+        localSdkRegistry = StubLocalSdkRegistry()
         appOwnedSdkRegistry = StubAppOwnedSdkInterfaceRegistry()
-        controller = LocalController(SDK_PACKAGE_NAME, locallyLoadedSdks, appOwnedSdkRegistry)
+        controller = LocalController(SDK_PACKAGE_NAME, localSdkRegistry, appOwnedSdkRegistry)
     }
 
     @Test
-    fun getSandboxedSdks_returnsResultsFromLocallyLoadedSdks() {
+    fun loadSdk_whenSdkRegistryReturnsResult_returnResultFromSdkRegistry() {
+        val expectedResult = SandboxedSdkCompat(Binder())
+        localSdkRegistry.loadSdkResult = expectedResult
+
+        val sdkParams = Bundle()
+        val callback = StubLoadSdkCallback()
+
+        controller.loadSdk(SDK_PACKAGE_NAME, sdkParams, Runnable::run, callback)
+
+        assertThat(callback.lastResult).isEqualTo(expectedResult)
+        assertThat(callback.lastError).isNull()
+
+        assertThat(localSdkRegistry.lastLoadSdkName).isEqualTo(SDK_PACKAGE_NAME)
+        assertThat(localSdkRegistry.lastLoadSdkParams).isSameInstanceAs(sdkParams)
+    }
+
+    @Test
+    fun loadSdk_whenSdkRegistryThrowsException_rethrowsExceptionFromSdkRegistry() {
+        val expectedError = LoadSdkCompatException(RuntimeException(), Bundle())
+        localSdkRegistry.loadSdkError = expectedError
+
+        val callback = StubLoadSdkCallback()
+
+        controller.loadSdk(SDK_PACKAGE_NAME, Bundle(), Runnable::run, callback)
+
+        assertThat(callback.lastError).isEqualTo(expectedError)
+        assertThat(callback.lastResult).isNull()
+    }
+
+    @Test
+    fun getSandboxedSdks_returnsResultsFromLocalSdkRegistry() {
         val sandboxedSdk = SandboxedSdkCompat(Binder())
-        locallyLoadedSdks.put(
-            "sdk", LocallyLoadedSdks.Entry(
-                sdkProvider = NoOpSdkProvider(),
-                sdk = sandboxedSdk
-            )
-        )
+        localSdkRegistry.getLoadedSdksResult = listOf(sandboxedSdk)
 
         val result = controller.getSandboxedSdks()
         assertThat(result).containsExactly(sandboxedSdk)
@@ -98,7 +124,7 @@
 
         val anotherSdkController = LocalController(
             "LocalControllerTest.anotherSdk",
-            locallyLoadedSdks,
+            localSdkRegistry,
             appOwnedSdkRegistry
         )
         val anotherSdkHandler = object : SdkSandboxActivityHandlerCompat {
@@ -130,14 +156,50 @@
         assertThat(registeredHandler).isNull()
     }
 
-    private class NoOpSdkProvider : LocalSdkProvider(Any()) {
-        override fun onLoadSdk(params: Bundle): SandboxedSdkCompat {
+    private class StubLocalSdkRegistry : SdkRegistry {
+
+        var getLoadedSdksResult: List<SandboxedSdkCompat> = emptyList()
+
+        var loadSdkResult: SandboxedSdkCompat? = null
+        var loadSdkError: LoadSdkCompatException? = null
+
+        var lastLoadSdkName: String? = null
+        var lastLoadSdkParams: Bundle? = null
+
+        override fun isResponsibleFor(sdkName: String): Boolean {
             throw IllegalStateException("Unexpected call")
         }
 
-        override fun beforeUnloadSdk() {
+        override fun loadSdk(sdkName: String, params: Bundle): SandboxedSdkCompat {
+            lastLoadSdkName = sdkName
+            lastLoadSdkParams = params
+
+            if (loadSdkError != null) {
+                throw loadSdkError!!
+            }
+
+            return loadSdkResult!!
+        }
+
+        override fun unloadSdk(sdkName: String) {
             throw IllegalStateException("Unexpected call")
         }
+
+        override fun getLoadedSdks(): List<SandboxedSdkCompat> = getLoadedSdksResult
+    }
+
+    private class StubLoadSdkCallback : SdkSandboxControllerCompat.LoadSdkCallback {
+
+        var lastResult: SandboxedSdkCompat? = null
+        var lastError: LoadSdkCompatException? = null
+
+        override fun onResult(result: SandboxedSdkCompat) {
+            lastResult = result
+        }
+
+        override fun onError(error: LoadSdkCompatException) {
+            lastError = error
+        }
     }
 
     private class StubAppOwnedSdkInterfaceRegistry : AppOwnedSdkRegistry {
diff --git a/privacysandbox/sdkruntime/sdkruntime-client/src/androidTest/java/androidx/privacysandbox/sdkruntime/client/controller/impl/LocalSdkRegistryTest.kt b/privacysandbox/sdkruntime/sdkruntime-client/src/androidTest/java/androidx/privacysandbox/sdkruntime/client/controller/impl/LocalSdkRegistryTest.kt
new file mode 100644
index 0000000..b386ed3
--- /dev/null
+++ b/privacysandbox/sdkruntime/sdkruntime-client/src/androidTest/java/androidx/privacysandbox/sdkruntime/client/controller/impl/LocalSdkRegistryTest.kt
@@ -0,0 +1,199 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.privacysandbox.sdkruntime.client.controller.impl
+
+import android.content.Context
+import android.os.Bundle
+import androidx.privacysandbox.sdkruntime.client.TestSdkConfigs
+import androidx.privacysandbox.sdkruntime.client.activity.LocalSdkActivityHandlerRegistry
+import androidx.privacysandbox.sdkruntime.client.loader.CatchingSdkActivityHandler
+import androidx.privacysandbox.sdkruntime.client.loader.asTestSdk
+import androidx.privacysandbox.sdkruntime.client.loader.extractSdkProviderFieldValue
+import androidx.privacysandbox.sdkruntime.core.LoadSdkCompatException
+import androidx.privacysandbox.sdkruntime.core.SandboxedSdkInfo
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.google.common.truth.Truth.assertThat
+import org.junit.Assert
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class LocalSdkRegistryTest {
+
+    private lateinit var context: Context
+    private lateinit var localSdkRegistry: LocalSdkRegistry
+
+    @Before
+    fun setUp() {
+        context = ApplicationProvider.getApplicationContext()
+        localSdkRegistry = LocalSdkRegistry.create(context, LocalAppOwnedSdkRegistry())
+    }
+
+    @Test
+    fun isResponsibleFor_LocalSdk_returnsTrue() {
+        val result = localSdkRegistry.isResponsibleFor(TestSdkConfigs.CURRENT.packageName)
+        assertThat(result).isTrue()
+    }
+
+    @Test
+    fun isResponsibleFor_NonLocalSdk_returnsFalse() {
+        val result = localSdkRegistry.isResponsibleFor("non-local-sdk")
+        assertThat(result).isFalse()
+    }
+
+    @Test
+    fun loadSdk_whenLocalSdkExists_returnsLocallyLoadedSdk() {
+        val result = localSdkRegistry.loadSdk(
+            TestSdkConfigs.CURRENT.packageName,
+            Bundle()
+        )
+
+        assertThat(result.getInterface()!!.javaClass.classLoader)
+            .isNotSameInstanceAs(localSdkRegistry.javaClass.classLoader)
+
+        assertThat(result.getSdkInfo())
+            .isEqualTo(
+                SandboxedSdkInfo(
+                    name = TestSdkConfigs.CURRENT.packageName,
+                    version = 42
+                )
+            )
+
+        assertThat(localSdkRegistry.getLoadedSdks()).containsExactly(result)
+    }
+
+    @Test
+    fun loadSdk_whenLocalSdkExists_rethrowsExceptionFromLocallyLoadedSdk() {
+        val params = Bundle()
+        params.putBoolean("needFail", true)
+
+        val result = Assert.assertThrows(LoadSdkCompatException::class.java) {
+            localSdkRegistry.loadSdk(
+                TestSdkConfigs.CURRENT.packageName,
+                params
+            )
+        }
+
+        assertThat(result.extraInformation).isEqualTo(params)
+        assertThat(result.loadSdkErrorCode).isEqualTo(
+            LoadSdkCompatException.LOAD_SDK_SDK_DEFINED_ERROR
+        )
+        assertThat(localSdkRegistry.getLoadedSdks()).isEmpty()
+    }
+
+    @Test
+    fun loadSdk_whenLocalSdkFailedToLoad_throwsInternalErrorException() {
+        val result = Assert.assertThrows(LoadSdkCompatException::class.java) {
+            localSdkRegistry.loadSdk(
+                TestSdkConfigs.forSdkName("invalidEntryPoint").packageName,
+                Bundle()
+            )
+        }
+
+        assertThat(result.loadSdkErrorCode).isEqualTo(
+            LoadSdkCompatException.LOAD_SDK_INTERNAL_ERROR
+        )
+        assertThat(result.message).isEqualTo("Failed to instantiate local SDK")
+        assertThat(localSdkRegistry.getLoadedSdks()).isEmpty()
+    }
+
+    @Test
+    fun loadSdk_whenSdkAlreadyLoaded_throwsSdkAlreadyLoadedException() {
+        val sdkName = TestSdkConfigs.CURRENT.packageName
+        val firstTimeLoadedSdk = localSdkRegistry.loadSdk(sdkName, Bundle())
+
+        val result = Assert.assertThrows(LoadSdkCompatException::class.java) {
+            localSdkRegistry.loadSdk(sdkName, Bundle())
+        }
+
+        assertThat(result.loadSdkErrorCode).isEqualTo(
+            LoadSdkCompatException.LOAD_SDK_ALREADY_LOADED
+        )
+        assertThat(localSdkRegistry.getLoadedSdks()).containsExactly(firstTimeLoadedSdk)
+    }
+
+    @Test
+    fun loadSdk_whenNoLocalSdkExists_throwsSdkNotFoundException() {
+        val result = Assert.assertThrows(LoadSdkCompatException::class.java) {
+            localSdkRegistry.loadSdk(
+                "sdk-doesnt-exist",
+                Bundle()
+            )
+        }
+
+        assertThat(result.loadSdkErrorCode).isEqualTo(LoadSdkCompatException.LOAD_SDK_NOT_FOUND)
+        assertThat(localSdkRegistry.getLoadedSdks()).isEmpty()
+    }
+
+    @Test
+    fun loadSdk_afterUnloading_loadSdkAgain() {
+        val sdkName = TestSdkConfigs.CURRENT.packageName
+        val sdkToUnload = localSdkRegistry.loadSdk(sdkName, Bundle())
+
+        localSdkRegistry.unloadSdk(sdkName)
+        val reloadedSdk = localSdkRegistry.loadSdk(sdkName, Bundle())
+
+        assertThat(localSdkRegistry.getLoadedSdks())
+            .containsExactly(reloadedSdk)
+        assertThat(reloadedSdk.getInterface())
+            .isNotEqualTo(sdkToUnload.getInterface())
+    }
+
+    @Test
+    fun unloadSdk_whenLocalSdkLoaded_unloadLocallyLoadedSdk() {
+        val sdkName = TestSdkConfigs.CURRENT.packageName
+        localSdkRegistry.loadSdk(sdkName, Bundle())
+        val sdkProvider = localSdkRegistry.getLoadedSdkProvider(sdkName)!!
+
+        localSdkRegistry.unloadSdk(sdkName)
+
+        val isBeforeUnloadSdkCalled = sdkProvider.extractSdkProviderFieldValue<Boolean>(
+            fieldName = "isBeforeUnloadSdkCalled"
+        )
+        assertThat(isBeforeUnloadSdkCalled).isTrue()
+        assertThat(localSdkRegistry.getLoadedSdks()).isEmpty()
+    }
+
+    @Test
+    fun unloadSdk_whenNoLocalSdkLoaded_doesntThrow() {
+        localSdkRegistry.unloadSdk(TestSdkConfigs.CURRENT.packageName)
+    }
+
+    @Test
+    fun unloadSdk_unregisterActivityHandlers() {
+        val packageName = TestSdkConfigs.CURRENT.packageName
+        val localSdk = localSdkRegistry.loadSdk(
+            packageName,
+            Bundle()
+        )
+
+        val testSdk = localSdk.asTestSdk()
+        val token = testSdk.registerSdkSandboxActivityHandler(CatchingSdkActivityHandler())
+
+        val registeredBefore = LocalSdkActivityHandlerRegistry.isRegistered(token)
+        assertThat(registeredBefore).isTrue()
+
+        localSdkRegistry.unloadSdk(packageName)
+
+        val registeredAfter = LocalSdkActivityHandlerRegistry.isRegistered(token)
+        assertThat(registeredAfter).isFalse()
+    }
+}
diff --git a/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/SdkSandboxManagerCompat.kt b/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/SdkSandboxManagerCompat.kt
index e04bd1ff..1de3a50 100644
--- a/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/SdkSandboxManagerCompat.kt
+++ b/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/SdkSandboxManagerCompat.kt
@@ -28,19 +28,15 @@
 import androidx.annotation.RequiresApi
 import androidx.core.os.BuildCompat
 import androidx.core.os.asOutcomeReceiver
-import androidx.privacysandbox.sdkruntime.client.activity.LocalSdkActivityHandlerRegistry
 import androidx.privacysandbox.sdkruntime.client.activity.LocalSdkActivityStarter
-import androidx.privacysandbox.sdkruntime.client.config.LocalSdkConfigsHolder
 import androidx.privacysandbox.sdkruntime.client.controller.AppOwnedSdkRegistry
-import androidx.privacysandbox.sdkruntime.client.controller.LocalControllerFactory
-import androidx.privacysandbox.sdkruntime.client.controller.LocallyLoadedSdks
+import androidx.privacysandbox.sdkruntime.client.controller.SdkRegistry
 import androidx.privacysandbox.sdkruntime.client.controller.impl.LocalAppOwnedSdkRegistry
+import androidx.privacysandbox.sdkruntime.client.controller.impl.LocalSdkRegistry
 import androidx.privacysandbox.sdkruntime.client.controller.impl.PlatformAppOwnedSdkRegistry
-import androidx.privacysandbox.sdkruntime.client.loader.SdkLoader
 import androidx.privacysandbox.sdkruntime.core.AdServicesInfo
 import androidx.privacysandbox.sdkruntime.core.AppOwnedSdkSandboxInterfaceCompat
 import androidx.privacysandbox.sdkruntime.core.LoadSdkCompatException
-import androidx.privacysandbox.sdkruntime.core.LoadSdkCompatException.Companion.LOAD_SDK_ALREADY_LOADED
 import androidx.privacysandbox.sdkruntime.core.LoadSdkCompatException.Companion.LOAD_SDK_NOT_FOUND
 import androidx.privacysandbox.sdkruntime.core.LoadSdkCompatException.Companion.toLoadCompatSdkException
 import androidx.privacysandbox.sdkruntime.core.SandboxedSdkCompat
@@ -94,10 +90,8 @@
  */
 class SdkSandboxManagerCompat private constructor(
     private val platformApi: PlatformApi,
-    private val configHolder: LocalSdkConfigsHolder,
-    private val localLocallyLoadedSdks: LocallyLoadedSdks,
-    private val appOwnedSdkRegistry: AppOwnedSdkRegistry,
-    private val sdkLoader: SdkLoader
+    private val localSdkRegistry: SdkRegistry,
+    private val appOwnedSdkRegistry: AppOwnedSdkRegistry
 ) {
     /**
      * Load SDK in a SDK sandbox java process or locally.
@@ -128,23 +122,10 @@
         sdkName: String,
         params: Bundle
     ): SandboxedSdkCompat {
-        if (localLocallyLoadedSdks.isLoaded(sdkName)) {
-            throw LoadSdkCompatException(LOAD_SDK_ALREADY_LOADED, "$sdkName already loaded")
+        val isLocalSdk = localSdkRegistry.isResponsibleFor(sdkName)
+        if (isLocalSdk) {
+            return localSdkRegistry.loadSdk(sdkName, params)
         }
-
-        val sdkConfig = configHolder.getSdkConfig(sdkName)
-        if (sdkConfig != null) {
-            val sdkProvider = sdkLoader.loadSdk(sdkConfig)
-            val sandboxedSdkCompat = sdkProvider.onLoadSdk(params)
-            localLocallyLoadedSdks.put(
-                sdkName, LocallyLoadedSdks.Entry(
-                    sdkProvider = sdkProvider,
-                    sdk = sandboxedSdkCompat
-                )
-            )
-            return sandboxedSdkCompat
-        }
-
         return platformApi.loadSdk(sdkName, params)
     }
 
@@ -158,12 +139,11 @@
      * @see [SdkSandboxManager.unloadSdk]
      */
     fun unloadSdk(sdkName: String) {
-        val localEntry = localLocallyLoadedSdks.remove(sdkName)
-        if (localEntry == null) {
-            platformApi.unloadSdk(sdkName)
+        val isLocalSdk = localSdkRegistry.isResponsibleFor(sdkName)
+        if (isLocalSdk) {
+            localSdkRegistry.unloadSdk(sdkName)
         } else {
-            localEntry.sdkProvider.beforeUnloadSdk()
-            LocalSdkActivityHandlerRegistry.unregisterAllActivityHandlersForSdk(sdkName)
+            platformApi.unloadSdk(sdkName)
         }
     }
 
@@ -210,7 +190,7 @@
      */
     fun getSandboxedSdks(): List<SandboxedSdkCompat> {
         val platformResult = platformApi.getSandboxedSdks()
-        val localResult = localLocallyLoadedSdks.getLoadedSdks()
+        val localResult = localSdkRegistry.getLoadedSdks()
         return platformResult + localResult
     }
 
@@ -265,10 +245,6 @@
         platformApi.startSdkSandboxActivity(fromActivity, sdkActivityToken)
     }
 
-    @TestOnly
-    internal fun getLocallyLoadedSdk(sdkName: String): LocallyLoadedSdks.Entry? =
-        localLocallyLoadedSdks.get(sdkName)
-
     private interface PlatformApi {
         @DoNotInline
         suspend fun loadSdk(sdkName: String, params: Bundle): SandboxedSdkCompat
@@ -426,18 +402,13 @@
                 val reference = sInstances[context]
                 var instance = reference?.get()
                 if (instance == null) {
-                    val configHolder = LocalSdkConfigsHolder.load(context)
-                    val localSdks = LocallyLoadedSdks()
                     val appOwnedSdkRegistry = AppOwnedSdkRegistryFactory.create(context)
-                    val controllerFactory = LocalControllerFactory(localSdks, appOwnedSdkRegistry)
-                    val sdkLoader = SdkLoader.create(context, controllerFactory)
+                    val localSdkRegistry = LocalSdkRegistry.create(context, appOwnedSdkRegistry)
                     val platformApi = PlatformApiFactory.create(context)
                     instance = SdkSandboxManagerCompat(
                         platformApi,
-                        configHolder,
-                        localSdks,
-                        appOwnedSdkRegistry,
-                        sdkLoader
+                        localSdkRegistry,
+                        appOwnedSdkRegistry
                     )
                     sInstances[context] = WeakReference(instance)
                 }
diff --git a/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/LocalController.kt b/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/LocalController.kt
index a57674c..5780689 100644
--- a/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/LocalController.kt
+++ b/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/LocalController.kt
@@ -31,7 +31,7 @@
  */
 internal class LocalController(
     private val sdkPackageName: String,
-    private val locallyLoadedSdks: LocallyLoadedSdks,
+    private val localSdkRegistry: SdkRegistry,
     private val appOwnedSdkRegistry: AppOwnedSdkRegistry
 ) : SdkSandboxControllerCompat.SandboxControllerImpl {
 
@@ -41,18 +41,20 @@
         executor: Executor,
         callback: SdkSandboxControllerCompat.LoadSdkCallback
     ) {
-        executor.execute {
-            callback.onError(
-                LoadSdkCompatException(
-                    LoadSdkCompatException.LOAD_SDK_INTERNAL_ERROR,
-                    "Shouldn't be called"
-                )
-            )
+        try {
+            val result = localSdkRegistry.loadSdk(sdkName, params)
+            executor.execute {
+                callback.onResult(result)
+            }
+        } catch (ex: LoadSdkCompatException) {
+            executor.execute {
+                callback.onError(ex)
+            }
         }
     }
 
     override fun getSandboxedSdks(): List<SandboxedSdkCompat> {
-        return locallyLoadedSdks.getLoadedSdks()
+        return localSdkRegistry.getLoadedSdks()
     }
 
     override fun getAppOwnedSdkSandboxInterfaces(): List<AppOwnedSdkSandboxInterfaceCompat> =
diff --git a/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/LocalControllerFactory.kt b/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/LocalControllerFactory.kt
index e6d1f0c..37932c8 100644
--- a/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/LocalControllerFactory.kt
+++ b/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/LocalControllerFactory.kt
@@ -24,12 +24,12 @@
  * Create [LocalController] instance for specific sdk.
  */
 internal class LocalControllerFactory(
-    private val locallyLoadedSdks: LocallyLoadedSdks,
+    private val localSdkRegistry: SdkRegistry,
     private val appOwnedSdkRegistry: AppOwnedSdkRegistry
 ) : SdkLoader.ControllerFactory {
     override fun createControllerFor(
         sdkConfig: LocalSdkConfig
     ): SdkSandboxControllerCompat.SandboxControllerImpl {
-        return LocalController(sdkConfig.packageName, locallyLoadedSdks, appOwnedSdkRegistry)
+        return LocalController(sdkConfig.packageName, localSdkRegistry, appOwnedSdkRegistry)
     }
 }
diff --git a/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/LocallyLoadedSdks.kt b/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/LocallyLoadedSdks.kt
deleted file mode 100644
index 73c3248..0000000
--- a/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/LocallyLoadedSdks.kt
+++ /dev/null
@@ -1,56 +0,0 @@
-/*
- * Copyright 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package androidx.privacysandbox.sdkruntime.client.controller
-
-import androidx.privacysandbox.sdkruntime.client.loader.LocalSdkProvider
-import androidx.privacysandbox.sdkruntime.core.SandboxedSdkCompat
-import org.jetbrains.annotations.TestOnly
-
-/**
- * Represents list of locally loaded SDKs.
- * Shared between:
- * 1) [androidx.privacysandbox.sdkruntime.client.SdkSandboxManagerCompat]
- * 2) [androidx.privacysandbox.sdkruntime.core.controller.SdkSandboxControllerCompat]
- */
-internal class LocallyLoadedSdks {
-
-    private val sdks = HashMap<String, Entry>()
-
-    fun isLoaded(sdkName: String): Boolean {
-        return sdks.containsKey(sdkName)
-    }
-
-    fun put(sdkName: String, entry: Entry) {
-        sdks[sdkName] = entry
-    }
-
-    @TestOnly
-    fun get(sdkName: String): Entry? = sdks[sdkName]
-
-    fun remove(sdkName: String): Entry? {
-        return sdks.remove(sdkName)
-    }
-
-    fun getLoadedSdks(): List<SandboxedSdkCompat> {
-        return sdks.values.map { it.sdk }
-    }
-
-    data class Entry(
-        val sdkProvider: LocalSdkProvider,
-        val sdk: SandboxedSdkCompat
-    )
-}
diff --git a/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/SdkRegistry.kt b/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/SdkRegistry.kt
new file mode 100644
index 0000000..eeed123
--- /dev/null
+++ b/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/SdkRegistry.kt
@@ -0,0 +1,60 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.privacysandbox.sdkruntime.client.controller
+
+import android.os.Bundle
+import androidx.privacysandbox.sdkruntime.core.LoadSdkCompatException
+import androidx.privacysandbox.sdkruntime.core.SandboxedSdkCompat
+
+/**
+ * Responsible for lifecycle of particular SDKs (local, sandbox, test, etc).
+ */
+internal interface SdkRegistry {
+
+    /**
+     * Checks if SDK could be loaded / unloaded by this SdkRegistry.
+     *
+     * @param sdkName name of the SDK to be loaded / unloaded.
+     * @return true if SDK could be loaded / unloaded by this SdkRegistry or false otherwise
+     */
+    fun isResponsibleFor(sdkName: String): Boolean
+
+    /**
+     * Loads SDK.
+     *
+     * @param sdkName name of the SDK to be loaded.
+     * @param params additional parameters to be passed to the SDK in the form of a [Bundle]
+     *  as agreed between the client and the SDK.
+     * @return [SandboxedSdkCompat] from SDK on a successful run.
+     * @throws [LoadSdkCompatException] on fail or when SdkRegistry not responsible for this SDK.
+     */
+    fun loadSdk(sdkName: String, params: Bundle): SandboxedSdkCompat
+
+    /**
+     * Unloads an SDK that has been previously loaded by SdkRegistry.
+     *
+     * @param sdkName name of the SDK to be unloaded.
+     */
+    fun unloadSdk(sdkName: String)
+
+    /**
+     * Fetches information about Sdks that are loaded by this SdkRegistry.
+     *
+     * @return List of [SandboxedSdkCompat] containing all currently loaded sdks
+     */
+    fun getLoadedSdks(): List<SandboxedSdkCompat>
+}
diff --git a/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/impl/LocalSdkRegistry.kt b/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/impl/LocalSdkRegistry.kt
new file mode 100644
index 0000000..58ee51d
--- /dev/null
+++ b/privacysandbox/sdkruntime/sdkruntime-client/src/main/java/androidx/privacysandbox/sdkruntime/client/controller/impl/LocalSdkRegistry.kt
@@ -0,0 +1,134 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.privacysandbox.sdkruntime.client.controller.impl
+
+import android.content.Context
+import android.os.Bundle
+import android.util.Log
+import androidx.privacysandbox.sdkruntime.client.activity.LocalSdkActivityHandlerRegistry
+import androidx.privacysandbox.sdkruntime.client.config.LocalSdkConfigsHolder
+import androidx.privacysandbox.sdkruntime.client.controller.AppOwnedSdkRegistry
+import androidx.privacysandbox.sdkruntime.client.controller.LocalControllerFactory
+import androidx.privacysandbox.sdkruntime.client.controller.SdkRegistry
+import androidx.privacysandbox.sdkruntime.client.loader.LocalSdkProvider
+import androidx.privacysandbox.sdkruntime.client.loader.SdkLoader
+import androidx.privacysandbox.sdkruntime.core.LoadSdkCompatException
+import androidx.privacysandbox.sdkruntime.core.SandboxedSdkCompat
+import org.jetbrains.annotations.TestOnly
+
+/**
+ * Responsible for lifecycle of SDKs bundled with app.
+ * Shared between:
+ * 1) [androidx.privacysandbox.sdkruntime.client.SdkSandboxManagerCompat]
+ * 2) [androidx.privacysandbox.sdkruntime.core.controller.SdkSandboxControllerCompat]
+ */
+internal class LocalSdkRegistry(
+    private val configHolder: LocalSdkConfigsHolder,
+) : SdkRegistry {
+    private lateinit var sdkLoader: SdkLoader
+
+    private val sdks = HashMap<String, Entry>()
+
+    override fun isResponsibleFor(sdkName: String): Boolean {
+        return configHolder.getSdkConfig(sdkName) != null
+    }
+
+    override fun loadSdk(sdkName: String, params: Bundle): SandboxedSdkCompat {
+        val sdkConfig = configHolder.getSdkConfig(sdkName)
+        if (sdkConfig == null) {
+            throw LoadSdkCompatException(
+                LoadSdkCompatException.LOAD_SDK_NOT_FOUND,
+                "$sdkName not bundled with app"
+            )
+        }
+
+        synchronized(sdks) {
+            if (sdks.containsKey(sdkName)) {
+                throw LoadSdkCompatException(
+                    LoadSdkCompatException.LOAD_SDK_ALREADY_LOADED,
+                    "$sdkName already loaded"
+                )
+            }
+
+            val sdkProvider = sdkLoader.loadSdk(sdkConfig)
+            val sandboxedSdkCompat = sdkProvider.onLoadSdk(params)
+            sdks.put(
+                sdkName, Entry(
+                    sdkProvider = sdkProvider,
+                    sdk = sandboxedSdkCompat
+                )
+            )
+            return sandboxedSdkCompat
+        }
+    }
+
+    override fun unloadSdk(sdkName: String) {
+        val loadedEntry = synchronized(sdks) {
+            sdks.remove(sdkName)
+        }
+        if (loadedEntry == null) {
+            Log.w(LOG_TAG, "Unloading SDK that is not loaded - $sdkName")
+            return
+        }
+
+        loadedEntry.sdkProvider.beforeUnloadSdk()
+        LocalSdkActivityHandlerRegistry.unregisterAllActivityHandlersForSdk(sdkName)
+    }
+
+    override fun getLoadedSdks(): List<SandboxedSdkCompat> = synchronized(sdks) {
+        return sdks.values.map { it.sdk }
+    }
+
+    @TestOnly
+    fun getLoadedSdkProvider(sdkName: String): LocalSdkProvider? = synchronized(sdks) {
+        return sdks[sdkName]?.sdkProvider
+    }
+
+    private data class Entry(
+        val sdkProvider: LocalSdkProvider,
+        val sdk: SandboxedSdkCompat
+    )
+
+    companion object {
+        const val LOG_TAG = "LocalSdkRegistry"
+
+        /**
+         * Create and initialize all required components for loading SDKs bundled with app.
+         *
+         * @param context App context
+         * @param appOwnedSdkRegistry AppOwnedSdkRegistry for [LocalControllerFactory]
+         * @return LocalSdkRegistry that could load SDKs bundled with app.
+         */
+        fun create(
+            context: Context,
+            appOwnedSdkRegistry: AppOwnedSdkRegistry
+        ): LocalSdkRegistry {
+            val configHolder = LocalSdkConfigsHolder.load(context)
+
+            val localSdkRegistry = LocalSdkRegistry(configHolder)
+            localSdkRegistry.sdkLoader = SdkLoader.create(
+                context,
+                LocalControllerFactory(
+                    localSdkRegistry,
+                    appOwnedSdkRegistry
+                )
+            )
+
+            return localSdkRegistry
+        }
+    }
+}