Merge "Remove unnnecessary project dependencies in room projects" into androidx-main
diff --git a/appfunctions/appfunctions-common/src/main/java/androidx/appfunctions/ExecuteAppFunctionResponse.kt b/appfunctions/appfunctions-common/src/main/java/androidx/appfunctions/ExecuteAppFunctionResponse.kt
index e4c0d22..a2dcb52 100644
--- a/appfunctions/appfunctions-common/src/main/java/androidx/appfunctions/ExecuteAppFunctionResponse.kt
+++ b/appfunctions/appfunctions-common/src/main/java/androidx/appfunctions/ExecuteAppFunctionResponse.kt
@@ -64,5 +64,9 @@
     public class Error(
         /** The [AppFunctionException] when the function execution failed. */
         public val error: AppFunctionException,
-    ) : ExecuteAppFunctionResponse
+    ) : ExecuteAppFunctionResponse {
+        override fun toString(): String {
+            return "AppFunctionResponse.Error(error=$error)"
+        }
+    }
 }
diff --git a/appfunctions/appfunctions-compiler/src/main/java/androidx/appfunctions/compiler/core/AnnotatedAppFunctionSerializable.kt b/appfunctions/appfunctions-compiler/src/main/java/androidx/appfunctions/compiler/core/AnnotatedAppFunctionSerializable.kt
index d120b09..9589fae 100644
--- a/appfunctions/appfunctions-compiler/src/main/java/androidx/appfunctions/compiler/core/AnnotatedAppFunctionSerializable.kt
+++ b/appfunctions/appfunctions-compiler/src/main/java/androidx/appfunctions/compiler/core/AnnotatedAppFunctionSerializable.kt
@@ -23,6 +23,7 @@
 import com.google.devtools.ksp.getVisibility
 import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.google.devtools.ksp.symbol.KSFile
+import com.google.devtools.ksp.symbol.KSFunctionDeclaration
 import com.google.devtools.ksp.symbol.KSValueParameter
 import com.google.devtools.ksp.symbol.Visibility
 import com.squareup.kotlinpoet.ClassName
@@ -39,44 +40,105 @@
      * 2. **Property Parameters:** Only properties (declared with `val`) can be passed as parameters
      *    to the primary constructor.
      * 3. **Supported Types:** All properties must be of one of the supported types.
+     * 4. **Super Type Parameters:** All parameters in the primary constructor of a super type must
+     *    be present in the primary constructor of the subtype.
      *
      * @throws ProcessingException if the class does not adhere to the requirements
      */
     fun validate(): AnnotatedAppFunctionSerializable {
-        val primaryConstructor = appFunctionSerializableClass.primaryConstructor
-        if (primaryConstructor == null || primaryConstructor.parameters.isEmpty()) {
+        val validatedPrimaryConstructor =
+            appFunctionSerializableClass.validateSerializablePrimaryConstructor()
+        val superTypesWithSerializableAnnotation =
+            appFunctionSerializableClass.superTypes
+                .map { it.resolve().declaration as KSClassDeclaration }
+                .filter {
+                    it.annotations.findAnnotation(
+                        IntrospectionHelper.AppFunctionSerializableAnnotation.CLASS_NAME
+                    ) != null
+                }
+                .map { it }
+                .toSet()
+        val parametersToValidate =
+            validatedPrimaryConstructor.parameters.associateBy { it.name.toString() }.toMutableMap()
+
+        validateParameters(parametersToValidate, superTypesWithSerializableAnnotation)
+
+        return this
+    }
+
+    private fun validateParameters(
+        parametersToValidate: MutableMap<String, KSValueParameter>,
+        superTypesWithSerializableAnnotation: Set<KSClassDeclaration>
+    ) {
+        for (serializableSuperType in superTypesWithSerializableAnnotation) {
+            val superTypePrimaryConstructor =
+                serializableSuperType.validateSerializablePrimaryConstructor()
+
+            for (superTypeParameter in superTypePrimaryConstructor.parameters) {
+                // Parameter has now been visited
+                val parameterInSuperType =
+                    parametersToValidate.remove(superTypeParameter.name.toString())
+                if (parameterInSuperType == null) {
+                    throw ProcessingException(
+                        "App parameters in @AppFunctionSerializable " +
+                            "supertypes must be present in subtype",
+                        superTypeParameter
+                    )
+                }
+                validateSerializableParameter(parameterInSuperType)
+            }
+        }
+
+        // Validate the remaining parameters
+        if (parametersToValidate.isNotEmpty()) {
+            for ((_, parameterToValidate) in parametersToValidate) {
+                validateSerializableParameter(parameterToValidate)
+            }
+        }
+    }
+
+    private fun KSClassDeclaration.validateSerializablePrimaryConstructor(): KSFunctionDeclaration {
+        if (primaryConstructor == null) {
             throw ProcessingException(
-                "Classes annotated with AppFunctionSerializable must have a primary constructor with one or more properties.",
-                appFunctionSerializableClass
+                "Classes annotated with AppFunctionSerializable must have a primary constructor.",
+                this
+            )
+        }
+        val primaryConstructorDeclaration = checkNotNull(primaryConstructor)
+        if (primaryConstructorDeclaration.parameters.isEmpty()) {
+            throw ProcessingException(
+                "Classes annotated with AppFunctionSerializable must not have an empty " +
+                    "primary constructor.",
+                this
             )
         }
 
-        if (primaryConstructor.getVisibility() != Visibility.PUBLIC) {
+        if (primaryConstructorDeclaration.getVisibility() != Visibility.PUBLIC) {
             throw ProcessingException(
                 "The primary constructor of @AppFunctionSerializable must be public.",
                 appFunctionSerializableClass
             )
         }
+        return primaryConstructorDeclaration
+    }
 
-        for (ksValueParameter in primaryConstructor.parameters) {
-            if (!ksValueParameter.isVal) {
-                throw ProcessingException(
-                    "All parameters in @AppFunctionSerializable primary constructor must have getters",
-                    ksValueParameter
-                )
-            }
-
-            if (!isSupportedType(ksValueParameter.type)) {
-                throw ProcessingException(
-                    "AppFunctionSerializable properties must be one of the following types:\n" +
-                        SUPPORTED_TYPES_STRING +
-                        ", an @AppFunctionSerializable or a list of @AppFunctionSerializable\nbut found " +
-                        ksValueParameter.type.toTypeName(),
-                    ksValueParameter
-                )
-            }
+    private fun validateSerializableParameter(ksValueParameter: KSValueParameter) {
+        if (!ksValueParameter.isVal) {
+            throw ProcessingException(
+                "All parameters in @AppFunctionSerializable primary constructor must have getters",
+                ksValueParameter
+            )
         }
-        return this
+
+        if (!isSupportedType(ksValueParameter.type)) {
+            throw ProcessingException(
+                "AppFunctionSerializable properties must be one of the following types:\n" +
+                    SUPPORTED_TYPES_STRING +
+                    ", an @AppFunctionSerializable or a list of @AppFunctionSerializable\nbut found " +
+                    ksValueParameter.type.toTypeName(),
+                ksValueParameter
+            )
+        }
     }
 
     /** Returns the annotated class's properties as defined in its primary constructor. */
diff --git a/appfunctions/appfunctions-compiler/src/main/java/androidx/appfunctions/compiler/processors/AppFunctionInvokerProcessor.kt b/appfunctions/appfunctions-compiler/src/main/java/androidx/appfunctions/compiler/processors/AppFunctionInvokerProcessor.kt
index 6078f71..b82adc5 100644
--- a/appfunctions/appfunctions-compiler/src/main/java/androidx/appfunctions/compiler/processors/AppFunctionInvokerProcessor.kt
+++ b/appfunctions/appfunctions-compiler/src/main/java/androidx/appfunctions/compiler/processors/AppFunctionInvokerProcessor.kt
@@ -218,7 +218,7 @@
             .returns(Any::class.asTypeName().copy(nullable = true))
             .addCode(
                 buildCodeBlock {
-                    addStatement("val result = when (${functionIdentifierSpec.name}) {")
+                    addStatement("val result: Any? = when (${functionIdentifierSpec.name}) {")
                     indent()
                     for (appFunction in annotatedAppFunctions.appFunctionDeclarations) {
                         appendInvocationBranchStatement(
diff --git a/appfunctions/appfunctions-compiler/src/test/java/androidx/appfunctions/compiler/processors/AppFunctionSerializableProcessorTest.kt b/appfunctions/appfunctions-compiler/src/test/java/androidx/appfunctions/compiler/processors/AppFunctionSerializableProcessorTest.kt
index 6285c4f..c637ecd 100644
--- a/appfunctions/appfunctions-compiler/src/test/java/androidx/appfunctions/compiler/processors/AppFunctionSerializableProcessorTest.kt
+++ b/appfunctions/appfunctions-compiler/src/test/java/androidx/appfunctions/compiler/processors/AppFunctionSerializableProcessorTest.kt
@@ -67,6 +67,41 @@
     }
 
     @Test
+    fun testProcessor_validInheritedProperties_success() {
+        val report =
+            compilationTestHelper.compileAll(sourceFileNames = listOf("DerivedSerializable.KT"))
+
+        compilationTestHelper.assertSuccessWithSourceContent(
+            report = report,
+            expectGeneratedSourceFileName = "\$DerivedSerializableFactory.kt",
+            goldenFileName = "\$DerivedSerializableFactory.KT"
+        )
+        compilationTestHelper.assertSuccessWithSourceContent(
+            report = report,
+            expectGeneratedSourceFileName = "\$StringBaseSerializableFactory.kt",
+            goldenFileName = "\$StringBaseSerializableFactory.KT"
+        )
+        compilationTestHelper.assertSuccessWithSourceContent(
+            report = report,
+            expectGeneratedSourceFileName = "\$LongBaseSerializableFactory.kt",
+            goldenFileName = "\$LongBaseSerializableFactory.KT"
+        )
+    }
+
+    @Test
+    fun testProcessor_badlyInheritedProperties_success() {
+        val report =
+            compilationTestHelper.compileAll(
+                sourceFileNames = listOf("SubClassRenamedPropertySerializable.KT")
+            )
+
+        compilationTestHelper.assertErrorWithMessage(
+            report,
+            "App parameters in @AppFunctionSerializable supertypes must be present in subtype"
+        )
+    }
+
+    @Test
     fun testProcessor_differentPackageSerializableProperty_success() {
         val report =
             compilationTestHelper.compileAll(
diff --git a/appfunctions/appfunctions-compiler/src/test/test-data/input/DerivedSerializable.KT b/appfunctions/appfunctions-compiler/src/test/test-data/input/DerivedSerializable.KT
new file mode 100644
index 0000000..6de75de
--- /dev/null
+++ b/appfunctions/appfunctions-compiler/src/test/test-data/input/DerivedSerializable.KT
@@ -0,0 +1,17 @@
+package com.testdata
+
+import androidx.appfunctions.AppFunctionSerializable
+
+@AppFunctionSerializable open class StringBaseSerializable(open val name: String) {}
+
+@AppFunctionSerializable
+open class LongBaseSerializable(
+    open val age: Long,
+    override val name: String,
+) : StringBaseSerializable(name) {}
+
+@AppFunctionSerializable
+class DerivedSerializable(
+    override val name: String,
+    override val age: Long,
+) : LongBaseSerializable(age, name)
diff --git a/appfunctions/appfunctions-compiler/src/test/test-data/input/SubClassRenamedPropertySerializable.KT b/appfunctions/appfunctions-compiler/src/test/test-data/input/SubClassRenamedPropertySerializable.KT
new file mode 100644
index 0000000..01fe81e
--- /dev/null
+++ b/appfunctions/appfunctions-compiler/src/test/test-data/input/SubClassRenamedPropertySerializable.KT
@@ -0,0 +1,9 @@
+package com.testdata
+
+import androidx.appfunctions.AppFunctionSerializable
+
+@AppFunctionSerializable open class BaseSerializable(open val name: String) {}
+
+@AppFunctionSerializable
+class SubClassRenamedPropertySerializable(val anotherString: String) :
+    BaseSerializable(anotherString)
diff --git a/appfunctions/appfunctions-compiler/src/test/test-data/output/$AllPrimitiveInputFunctions_AppFunctionInvoker.KT b/appfunctions/appfunctions-compiler/src/test/test-data/output/$AllPrimitiveInputFunctions_AppFunctionInvoker.KT
index 22d5c14..3c243fd 100644
--- a/appfunctions/appfunctions-compiler/src/test/test-data/output/$AllPrimitiveInputFunctions_AppFunctionInvoker.KT
+++ b/appfunctions/appfunctions-compiler/src/test/test-data/output/$AllPrimitiveInputFunctions_AppFunctionInvoker.KT
@@ -44,7 +44,7 @@
     functionIdentifier: String,
     parameters: Map<String, Any?>,
   ): Any? {
-    val result = when (functionIdentifier) {
+    val result: Any? = when (functionIdentifier) {
       "com.testdata.AllPrimitiveInputFunctions#simpleFunctionInt" -> {
         ConfigurableAppFunctionFactory<AllPrimitiveInputFunctions>(
           appFunctionContext.context
diff --git a/appfunctions/appfunctions-compiler/src/test/test-data/output/$DerivedSerializableFactory.KT b/appfunctions/appfunctions-compiler/src/test/test-data/output/$DerivedSerializableFactory.KT
new file mode 100644
index 0000000..1ff1b9a
--- /dev/null
+++ b/appfunctions/appfunctions-compiler/src/test/test-data/output/$DerivedSerializableFactory.KT
@@ -0,0 +1,25 @@
+package com.testdata
+
+import androidx.appfunctions.AppFunctionData
+import androidx.appfunctions.`internal`.AppFunctionSerializableFactory
+import javax.`annotation`.processing.Generated
+
+@Generated("androidx.appfunctions.compiler.AppFunctionCompiler")
+public class `$DerivedSerializableFactory` : AppFunctionSerializableFactory<DerivedSerializable> {
+  override fun fromAppFunctionData(appFunctionData: AppFunctionData): DerivedSerializable {
+
+    val name = appFunctionData.getString("name")
+    val age = appFunctionData.getLong("age")
+
+    return DerivedSerializable(name, age)
+  }
+
+  override fun toAppFunctionData(derivedSerializable: DerivedSerializable): AppFunctionData {
+
+    val builder = AppFunctionData.Builder("com.testdata.DerivedSerializable")
+    builder.setString("name", derivedSerializable.name)
+    builder.setLong("age", derivedSerializable.age)
+
+    return builder.build()
+  }
+}
diff --git a/appfunctions/appfunctions-compiler/src/test/test-data/output/$LongBaseSerializableFactory.KT b/appfunctions/appfunctions-compiler/src/test/test-data/output/$LongBaseSerializableFactory.KT
new file mode 100644
index 0000000..79729b9
--- /dev/null
+++ b/appfunctions/appfunctions-compiler/src/test/test-data/output/$LongBaseSerializableFactory.KT
@@ -0,0 +1,25 @@
+package com.testdata
+
+import androidx.appfunctions.AppFunctionData
+import androidx.appfunctions.`internal`.AppFunctionSerializableFactory
+import javax.`annotation`.processing.Generated
+
+@Generated("androidx.appfunctions.compiler.AppFunctionCompiler")
+public class `$LongBaseSerializableFactory` : AppFunctionSerializableFactory<LongBaseSerializable> {
+  override fun fromAppFunctionData(appFunctionData: AppFunctionData): LongBaseSerializable {
+
+    val age = appFunctionData.getLong("age")
+    val name = appFunctionData.getString("name")
+
+    return LongBaseSerializable(age, name)
+  }
+
+  override fun toAppFunctionData(longBaseSerializable: LongBaseSerializable): AppFunctionData {
+
+    val builder = AppFunctionData.Builder("com.testdata.LongBaseSerializable")
+    builder.setLong("age", longBaseSerializable.age)
+    builder.setString("name", longBaseSerializable.name)
+
+    return builder.build()
+  }
+}
diff --git a/appfunctions/appfunctions-compiler/src/test/test-data/output/$SimpleFunction_AppFunctionInvoker.KT b/appfunctions/appfunctions-compiler/src/test/test-data/output/$SimpleFunction_AppFunctionInvoker.KT
index 19ea695..55b6e07 100644
--- a/appfunctions/appfunctions-compiler/src/test/test-data/output/$SimpleFunction_AppFunctionInvoker.KT
+++ b/appfunctions/appfunctions-compiler/src/test/test-data/output/$SimpleFunction_AppFunctionInvoker.KT
@@ -21,7 +21,7 @@
     functionIdentifier: String,
     parameters: Map<String, Any?>,
   ): Any? {
-    val result = when (functionIdentifier) {
+    val result: Any? = when (functionIdentifier) {
       "com.testdata.SimpleFunction#simpleFunction" -> {
         ConfigurableAppFunctionFactory<SimpleFunction>(
           appFunctionContext.context
diff --git a/appfunctions/appfunctions-compiler/src/test/test-data/output/$StringBaseSerializableFactory.KT b/appfunctions/appfunctions-compiler/src/test/test-data/output/$StringBaseSerializableFactory.KT
new file mode 100644
index 0000000..b8987a1
--- /dev/null
+++ b/appfunctions/appfunctions-compiler/src/test/test-data/output/$StringBaseSerializableFactory.KT
@@ -0,0 +1,24 @@
+package com.testdata
+
+import androidx.appfunctions.AppFunctionData
+import androidx.appfunctions.`internal`.AppFunctionSerializableFactory
+import javax.`annotation`.processing.Generated
+
+@Generated("androidx.appfunctions.compiler.AppFunctionCompiler")
+public class `$StringBaseSerializableFactory` :
+    AppFunctionSerializableFactory<StringBaseSerializable> {
+  override fun fromAppFunctionData(appFunctionData: AppFunctionData): StringBaseSerializable {
+
+    val name = appFunctionData.getString("name")
+
+    return StringBaseSerializable(name)
+  }
+
+  override fun toAppFunctionData(stringBaseSerializable: StringBaseSerializable): AppFunctionData {
+
+    val builder = AppFunctionData.Builder("com.testdata.StringBaseSerializable")
+    builder.setString("name", stringBaseSerializable.name)
+
+    return builder.build()
+  }
+}
diff --git a/appfunctions/appfunctions-runtime/api/current.txt b/appfunctions/appfunctions-runtime/api/current.txt
index 6434b84..17d48a8 100644
--- a/appfunctions/appfunctions-runtime/api/current.txt
+++ b/appfunctions/appfunctions-runtime/api/current.txt
@@ -1,7 +1,7 @@
 // Signature format: 4.0
 package androidx.appfunctions {
 
-  @kotlin.annotation.Retention(kotlin.annotation.AnnotationRetention.SOURCE) @kotlin.annotation.Target(allowedTargets=kotlin.annotation.AnnotationTarget.FUNCTION) public @interface AppFunction {
+  @kotlin.annotation.Retention(kotlin.annotation.AnnotationRetention.BINARY) @kotlin.annotation.Target(allowedTargets=kotlin.annotation.AnnotationTarget.FUNCTION) public @interface AppFunction {
     method public abstract boolean isEnabled() default true;
     property public abstract boolean isEnabled;
   }
diff --git a/appfunctions/appfunctions-runtime/api/restricted_current.txt b/appfunctions/appfunctions-runtime/api/restricted_current.txt
index 6434b84..17d48a8 100644
--- a/appfunctions/appfunctions-runtime/api/restricted_current.txt
+++ b/appfunctions/appfunctions-runtime/api/restricted_current.txt
@@ -1,7 +1,7 @@
 // Signature format: 4.0
 package androidx.appfunctions {
 
-  @kotlin.annotation.Retention(kotlin.annotation.AnnotationRetention.SOURCE) @kotlin.annotation.Target(allowedTargets=kotlin.annotation.AnnotationTarget.FUNCTION) public @interface AppFunction {
+  @kotlin.annotation.Retention(kotlin.annotation.AnnotationRetention.BINARY) @kotlin.annotation.Target(allowedTargets=kotlin.annotation.AnnotationTarget.FUNCTION) public @interface AppFunction {
     method public abstract boolean isEnabled() default true;
     property public abstract boolean isEnabled;
   }
diff --git a/appfunctions/appfunctions-runtime/build.gradle b/appfunctions/appfunctions-runtime/build.gradle
index 42fabd3..f73d92b 100644
--- a/appfunctions/appfunctions-runtime/build.gradle
+++ b/appfunctions/appfunctions-runtime/build.gradle
@@ -48,8 +48,8 @@
     testImplementation(libs.mockitoCore4)
 
     androidTestCompileOnly(project(":appfunctions:appfunctions-stubs"))
-    androidTestImplementation("androidx.appsearch:appsearch:1.1.0-alpha07")
-    androidTestImplementation("androidx.appsearch:appsearch-platform-storage:1.1.0-alpha07")
+    androidTestImplementation("androidx.appsearch:appsearch:1.1.0-beta01")
+    androidTestImplementation("androidx.appsearch:appsearch-platform-storage:1.1.0-beta01")
     androidTestImplementation(libs.junit)
     androidTestImplementation(libs.truth)
     androidTestImplementation(libs.kotlinCoroutinesTest)
diff --git a/appfunctions/appfunctions-runtime/src/androidTest/assets/app_function.xml b/appfunctions/appfunctions-runtime/src/androidTest/assets/app_functions.xml
similarity index 100%
rename from appfunctions/appfunctions-runtime/src/androidTest/assets/app_function.xml
rename to appfunctions/appfunctions-runtime/src/androidTest/assets/app_functions.xml
diff --git a/appfunctions/appfunctions-runtime/src/main/AndroidManifest.xml b/appfunctions/appfunctions-runtime/src/main/AndroidManifest.xml
index 56b769d..223cd2d 100644
--- a/appfunctions/appfunctions-runtime/src/main/AndroidManifest.xml
+++ b/appfunctions/appfunctions-runtime/src/main/AndroidManifest.xml
@@ -23,7 +23,7 @@
             android:exported="true">
             <property
                 android:name="android.app.appfunctions"
-                android:value="app_function.xml" />
+                android:value="app_functions.xml" />
             <property
                 android:name="android.app.appfunctions.schema"
                 android:value="app_functions_schema.xml" />
diff --git a/appfunctions/appfunctions-runtime/src/main/java/androidx/appfunctions/AppFunction.kt b/appfunctions/appfunctions-runtime/src/main/java/androidx/appfunctions/AppFunction.kt
index c6a4e1b..a700d19 100644
--- a/appfunctions/appfunctions-runtime/src/main/java/androidx/appfunctions/AppFunction.kt
+++ b/appfunctions/appfunctions-runtime/src/main/java/androidx/appfunctions/AppFunction.kt
@@ -62,7 +62,8 @@
  *
  * @see AppFunctionConfiguration.Builder.addEnclosingClassFactory
  */
-@Retention(AnnotationRetention.SOURCE)
+// Use BINARY here so that the annotation is kept around at the aggregation stage.
+@Retention(AnnotationRetention.BINARY)
 @Target(AnnotationTarget.FUNCTION)
 public annotation class AppFunction(
     /**
diff --git a/appfunctions/integration-tests/multi-modules-testapp/app/build.gradle b/appfunctions/integration-tests/multi-modules-testapp/app/build.gradle
index 49fd220..29fe921 100644
--- a/appfunctions/integration-tests/multi-modules-testapp/app/build.gradle
+++ b/appfunctions/integration-tests/multi-modules-testapp/app/build.gradle
@@ -25,17 +25,28 @@
     defaultConfig {
         applicationId = "androidx.appfunctions.integration.testapp"
         minSdk = 33
+        ksp {
+            arg("appfunctions:aggregateAppFunctions", "true")
+        }
     }
 }
 dependencies {
     api(libs.kotlinStdlib)
-    // AppFunctions libs
+    api(libs.kotlinCoroutinesAndroid)
     implementation(project(":appfunctions:appfunctions-common"))
     implementation(project(":appfunctions:appfunctions-runtime"))
-    // AppFunction compiler
     ksp(project(":appfunctions:appfunctions-compiler"))
-
-    // shared integration test module
     implementation(project(":appfunctions:integration-tests:multi-modules-testapp:shared-library"))
+    implementation('androidx.concurrent:concurrent-futures-ktx:1.1.0')
 
+    androidTestImplementation(libs.junit)
+    androidTestImplementation(libs.truth)
+    androidTestImplementation(libs.kotlinCoroutinesTest)
+    androidTestImplementation(libs.testRunner)
+    androidTestImplementation(libs.guavaAndroid)
+    androidTestImplementation(libs.testMonitor)
+    androidTestImplementation(libs.kotlinTest)
+    androidTestImplementation("androidx.appsearch:appsearch:1.1.0-beta01")
+    androidTestImplementation("androidx.appsearch:appsearch-platform-storage:1.1.0-beta01")
+    androidTestImplementation('androidx.concurrent:concurrent-futures-ktx:1.1.0')
 }
\ No newline at end of file
diff --git a/appfunctions/integration-tests/multi-modules-testapp/app/src/androidTest/java/androidx/appfunctions/integration/tests/AppSearchMetadataHelper.kt b/appfunctions/integration-tests/multi-modules-testapp/app/src/androidTest/java/androidx/appfunctions/integration/tests/AppSearchMetadataHelper.kt
new file mode 100644
index 0000000..86585978
--- /dev/null
+++ b/appfunctions/integration-tests/multi-modules-testapp/app/src/androidTest/java/androidx/appfunctions/integration/tests/AppSearchMetadataHelper.kt
@@ -0,0 +1,61 @@
+/*
+ * Copyright 2025 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.appfunctions.integration.tests
+
+import android.content.Context
+import androidx.appsearch.app.GlobalSearchSession
+import androidx.appsearch.app.SearchSpec
+import androidx.appsearch.platformstorage.PlatformStorage
+import androidx.concurrent.futures.await
+
+internal object AppSearchMetadataHelper {
+    /** Returns function IDs that belong to the given context's package. */
+    suspend fun collectSelfFunctionIds(context: Context): Set<String> {
+        val functionIds = mutableSetOf<String>()
+        createSearchSession(context).use { session ->
+            val searchResults =
+                session.search(
+                    "",
+                    SearchSpec.Builder()
+                        .addFilterNamespaces("app_functions_runtime")
+                        .addFilterPackageNames("android")
+                        .addFilterSchemas("AppFunctionRuntimeMetadata")
+                        .build(),
+                )
+            var nextPage = searchResults.nextPageAsync.await()
+            while (nextPage.isNotEmpty()) {
+                for (result in nextPage) {
+                    val packageName = result.genericDocument.getPropertyString("packageName")
+                    if (packageName != context.packageName) {
+                        continue
+                    }
+                    val functionId = result.genericDocument.getPropertyString("functionId")
+                    functionIds.add(checkNotNull(functionId))
+                }
+                nextPage = searchResults.nextPageAsync.await()
+            }
+        }
+        return functionIds
+    }
+
+    private suspend fun createSearchSession(context: Context): GlobalSearchSession {
+        return PlatformStorage.createGlobalSearchSessionAsync(
+                PlatformStorage.GlobalSearchContext.Builder(context).build()
+            )
+            .await()
+    }
+}
diff --git a/appfunctions/integration-tests/multi-modules-testapp/app/src/androidTest/java/androidx/appfunctions/integration/tests/IntegrationTest.kt b/appfunctions/integration-tests/multi-modules-testapp/app/src/androidTest/java/androidx/appfunctions/integration/tests/IntegrationTest.kt
new file mode 100644
index 0000000..8e7890b
--- /dev/null
+++ b/appfunctions/integration-tests/multi-modules-testapp/app/src/androidTest/java/androidx/appfunctions/integration/tests/IntegrationTest.kt
@@ -0,0 +1,119 @@
+/*
+ * Copyright 2025 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.appfunction.integration.tests
+
+import androidx.appfunctions.AppFunctionData
+import androidx.appfunctions.AppFunctionFunctionNotFoundException
+import androidx.appfunctions.AppFunctionInvalidArgumentException
+import androidx.appfunctions.AppFunctionManagerCompat
+import androidx.appfunctions.ExecuteAppFunctionRequest
+import androidx.appfunctions.ExecuteAppFunctionResponse
+import androidx.appfunctions.integration.tests.AppSearchMetadataHelper
+import androidx.appfunctions.integration.tests.TestUtil.doBlocking
+import androidx.appfunctions.integration.tests.TestUtil.retryAssert
+import androidx.test.filters.LargeTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.google.common.truth.Truth.assertThat
+import kotlin.test.assertIs
+import kotlinx.coroutines.delay
+import org.junit.Assume.assumeTrue
+import org.junit.Before
+import org.junit.Test
+
+@LargeTest
+class IntegrationTest {
+    private val targetContext = InstrumentationRegistry.getInstrumentation().targetContext
+    private val appFunctionManager = AppFunctionManagerCompat(targetContext)
+
+    @Before
+    fun setup() = doBlocking {
+        assumeTrue(appFunctionManager.isSupported())
+
+        awaitAppFunctionsIndexed(FUNCTION_IDS)
+    }
+
+    @Test
+    fun executeAppFunction_success() = doBlocking {
+        val response =
+            appFunctionManager.executeAppFunction(
+                request =
+                    ExecuteAppFunctionRequest(
+                        targetContext.packageName,
+                        "androidx.appfunctions.integration.testapp.TestFunctions#add",
+                        AppFunctionData.Builder("").setLong("num1", 1).setLong("num2", 2).build()
+                    )
+            )
+
+        val successResponse = assertIs<ExecuteAppFunctionResponse.Success>(response)
+        assertThat(
+                successResponse.returnValue.getLong(
+                    ExecuteAppFunctionResponse.Success.PROPERTY_RETURN_VALUE
+                )
+            )
+            .isEqualTo(3)
+    }
+
+    @Test
+    fun executeAppFunction_functionNotFound_fail() = doBlocking {
+        val response =
+            appFunctionManager.executeAppFunction(
+                request =
+                    ExecuteAppFunctionRequest(
+                        targetContext.packageName,
+                        "androidx.appfunctions.integration.testapp.TestFunctions#notExist",
+                        AppFunctionData.Builder("").build()
+                    )
+            )
+
+        val errorResponse = assertIs<ExecuteAppFunctionResponse.Error>(response)
+        assertThat(errorResponse.error)
+            .isInstanceOf(AppFunctionFunctionNotFoundException::class.java)
+    }
+
+    @Test
+    fun executeAppFunction_appThrows_fail() = doBlocking {
+        delay(5)
+        val response =
+            appFunctionManager.executeAppFunction(
+                request =
+                    ExecuteAppFunctionRequest(
+                        targetContext.packageName,
+                        "androidx.appfunctions.integration.testapp.TestFunctions#doThrow",
+                        AppFunctionData.Builder("").build()
+                    )
+            )
+
+        assertThat(response).isInstanceOf(ExecuteAppFunctionResponse.Error::class.java)
+        val errorResponse = response as ExecuteAppFunctionResponse.Error
+        assertThat(errorResponse.error)
+            .isInstanceOf(AppFunctionInvalidArgumentException::class.java)
+    }
+
+    private suspend fun awaitAppFunctionsIndexed(expectedFunctionIds: Set<String>) {
+        retryAssert {
+            val functionIds = AppSearchMetadataHelper.collectSelfFunctionIds(targetContext)
+            assertThat(functionIds).containsAtLeastElementsIn(expectedFunctionIds)
+        }
+    }
+
+    private companion object {
+        const val APP_FUNCTION_ID = "androidx.appfunctions.integration.testapp.TestFunctions#add"
+        const val DO_THROW_FUNCTION_ID =
+            "androidx.appfunctions.integration.testapp.TestFunctions#doThrow"
+        val FUNCTION_IDS = setOf(APP_FUNCTION_ID, DO_THROW_FUNCTION_ID)
+    }
+}
diff --git a/appfunctions/integration-tests/multi-modules-testapp/app/src/androidTest/java/androidx/appfunctions/integration/tests/TestUtil.kt b/appfunctions/integration-tests/multi-modules-testapp/app/src/androidTest/java/androidx/appfunctions/integration/tests/TestUtil.kt
new file mode 100644
index 0000000..aca1e29
--- /dev/null
+++ b/appfunctions/integration-tests/multi-modules-testapp/app/src/androidTest/java/androidx/appfunctions/integration/tests/TestUtil.kt
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2025 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.appfunctions.integration.tests
+
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.runBlocking
+
+internal object TestUtil {
+    fun doBlocking(block: suspend CoroutineScope.() -> Unit) = runBlocking(block = block)
+
+    fun interface ThrowRunnable {
+        @Throws(Throwable::class) suspend fun run()
+    }
+
+    /** Retries an assertion with a delay between attempts. */
+    @Throws(Throwable::class)
+    suspend fun retryAssert(runnable: ThrowRunnable) {
+        var lastError: Throwable? = null
+
+        for (attempt in 0 until RETRY_MAX_INTERVALS) {
+            try {
+                runnable.run()
+                return
+            } catch (e: Throwable) {
+                lastError = e
+                delay(RETRY_CHECK_INTERVAL_MILLIS)
+            }
+        }
+        throw lastError!!
+    }
+
+    private const val RETRY_CHECK_INTERVAL_MILLIS: Long = 500
+    private const val RETRY_MAX_INTERVALS: Long = 10
+}
diff --git a/appfunctions/integration-tests/multi-modules-testapp/app/src/main/java/androidx/appfunctions/integration/testapp/TestFunctions.kt b/appfunctions/integration-tests/multi-modules-testapp/app/src/main/java/androidx/appfunctions/integration/testapp/TestFunctions.kt
new file mode 100644
index 0000000..dbe8a32
--- /dev/null
+++ b/appfunctions/integration-tests/multi-modules-testapp/app/src/main/java/androidx/appfunctions/integration/testapp/TestFunctions.kt
@@ -0,0 +1,32 @@
+/*
+ * Copyright 2025 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.appfunctions.integration.testapp
+
+import androidx.appfunctions.AppFunction
+import androidx.appfunctions.AppFunctionContext
+import androidx.appfunctions.AppFunctionInvalidArgumentException
+
+@Suppress("UNUSED_PARAMETER")
+class TestFunctions {
+    @AppFunction
+    fun add(appFunctionContext: AppFunctionContext, num1: Long, num2: Long) = num1 + num2
+
+    @AppFunction
+    fun doThrow(appFunctionContext: AppFunctionContext) {
+        throw AppFunctionInvalidArgumentException("invalid")
+    }
+}
diff --git a/appfunctions/integration-tests/multi-modules-testapp/shared-library/build.gradle b/appfunctions/integration-tests/multi-modules-testapp/shared-library/build.gradle
index 3b656c2..35fcb96 100644
--- a/appfunctions/integration-tests/multi-modules-testapp/shared-library/build.gradle
+++ b/appfunctions/integration-tests/multi-modules-testapp/shared-library/build.gradle
@@ -31,9 +31,7 @@
 
 dependencies {
     api(libs.kotlinStdlib)
-    // AppFunctions libs
     implementation(project(":appfunctions:appfunctions-common"))
     implementation(project(":appfunctions:appfunctions-runtime"))
-    // AppFunction compiler
     ksp(project(":appfunctions:appfunctions-compiler"))
 }
\ No newline at end of file
diff --git a/compose/material3/material3/api/current.txt b/compose/material3/material3/api/current.txt
index d06a35b..c0a37e8 100644
--- a/compose/material3/material3/api/current.txt
+++ b/compose/material3/material3/api/current.txt
@@ -20,6 +20,8 @@
   public final class AlertDialogKt {
     method @Deprecated @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3Api @androidx.compose.runtime.Composable public static void AlertDialog(kotlin.jvm.functions.Function0<kotlin.Unit> onDismissRequest, optional androidx.compose.ui.Modifier modifier, optional androidx.compose.ui.window.DialogProperties properties, kotlin.jvm.functions.Function0<kotlin.Unit> content);
     method @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3Api @androidx.compose.runtime.Composable public static void BasicAlertDialog(kotlin.jvm.functions.Function0<kotlin.Unit> onDismissRequest, optional androidx.compose.ui.Modifier modifier, optional androidx.compose.ui.window.DialogProperties properties, kotlin.jvm.functions.Function0<kotlin.Unit> content);
+    method @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3ComponentOverrideApi public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.compose.material3.BasicAlertDialogOverride> getLocalBasicAlertDialogOverride();
+    property @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3ComponentOverrideApi public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.compose.material3.BasicAlertDialogOverride> LocalBasicAlertDialogOverride;
   }
 
   public final class AndroidAlertDialog_androidKt {
@@ -85,6 +87,21 @@
     method @androidx.compose.runtime.Composable public static void BadgedBox(kotlin.jvm.functions.Function1<? super androidx.compose.foundation.layout.BoxScope,kotlin.Unit> badge, optional androidx.compose.ui.Modifier modifier, kotlin.jvm.functions.Function1<? super androidx.compose.foundation.layout.BoxScope,kotlin.Unit> content);
   }
 
+  @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3ComponentOverrideApi public interface BasicAlertDialogOverride {
+    method @androidx.compose.runtime.Composable public void BasicAlertDialog(androidx.compose.material3.BasicAlertDialogOverrideScope);
+  }
+
+  @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3ComponentOverrideApi public final class BasicAlertDialogOverrideScope {
+    method public kotlin.jvm.functions.Function0<kotlin.Unit> getContent();
+    method public androidx.compose.ui.Modifier getModifier();
+    method public kotlin.jvm.functions.Function0<kotlin.Unit> getOnDismissRequest();
+    method public androidx.compose.ui.window.DialogProperties getProperties();
+    property public kotlin.jvm.functions.Function0<kotlin.Unit> content;
+    property public androidx.compose.ui.Modifier modifier;
+    property public kotlin.jvm.functions.Function0<kotlin.Unit> onDismissRequest;
+    property public androidx.compose.ui.window.DialogProperties properties;
+  }
+
   public final class BottomAppBarDefaults {
     method @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3Api @androidx.compose.runtime.Composable public androidx.compose.material3.BottomAppBarScrollBehavior exitAlwaysScrollBehavior(optional androidx.compose.material3.BottomAppBarState state, optional kotlin.jvm.functions.Function0<java.lang.Boolean> canScroll, optional androidx.compose.animation.core.AnimationSpec<java.lang.Float>? snapAnimationSpec, optional androidx.compose.animation.core.DecayAnimationSpec<java.lang.Float>? flingAnimationSpec);
     method @androidx.compose.runtime.Composable public long getBottomAppBarFabColor();
@@ -721,6 +738,11 @@
     property public abstract kotlin.ranges.IntRange yearRange;
   }
 
+  @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3ComponentOverrideApi public final class DefaultBasicAlertDialogOverride implements androidx.compose.material3.BasicAlertDialogOverride {
+    method @androidx.compose.runtime.Composable public void BasicAlertDialog(androidx.compose.material3.BasicAlertDialogOverrideScope);
+    field public static final androidx.compose.material3.DefaultBasicAlertDialogOverride INSTANCE;
+  }
+
   @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3ComponentOverrideApi public final class DefaultNavigationBarComponentOverride implements androidx.compose.material3.NavigationBarComponentOverride {
     method @androidx.compose.runtime.Composable public void NavigationBar(androidx.compose.material3.NavigationBarComponentOverrideContext);
     field public static final androidx.compose.material3.DefaultNavigationBarComponentOverride INSTANCE;
diff --git a/compose/material3/material3/api/restricted_current.txt b/compose/material3/material3/api/restricted_current.txt
index d06a35b..c0a37e8 100644
--- a/compose/material3/material3/api/restricted_current.txt
+++ b/compose/material3/material3/api/restricted_current.txt
@@ -20,6 +20,8 @@
   public final class AlertDialogKt {
     method @Deprecated @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3Api @androidx.compose.runtime.Composable public static void AlertDialog(kotlin.jvm.functions.Function0<kotlin.Unit> onDismissRequest, optional androidx.compose.ui.Modifier modifier, optional androidx.compose.ui.window.DialogProperties properties, kotlin.jvm.functions.Function0<kotlin.Unit> content);
     method @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3Api @androidx.compose.runtime.Composable public static void BasicAlertDialog(kotlin.jvm.functions.Function0<kotlin.Unit> onDismissRequest, optional androidx.compose.ui.Modifier modifier, optional androidx.compose.ui.window.DialogProperties properties, kotlin.jvm.functions.Function0<kotlin.Unit> content);
+    method @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3ComponentOverrideApi public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.compose.material3.BasicAlertDialogOverride> getLocalBasicAlertDialogOverride();
+    property @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3ComponentOverrideApi public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.compose.material3.BasicAlertDialogOverride> LocalBasicAlertDialogOverride;
   }
 
   public final class AndroidAlertDialog_androidKt {
@@ -85,6 +87,21 @@
     method @androidx.compose.runtime.Composable public static void BadgedBox(kotlin.jvm.functions.Function1<? super androidx.compose.foundation.layout.BoxScope,kotlin.Unit> badge, optional androidx.compose.ui.Modifier modifier, kotlin.jvm.functions.Function1<? super androidx.compose.foundation.layout.BoxScope,kotlin.Unit> content);
   }
 
+  @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3ComponentOverrideApi public interface BasicAlertDialogOverride {
+    method @androidx.compose.runtime.Composable public void BasicAlertDialog(androidx.compose.material3.BasicAlertDialogOverrideScope);
+  }
+
+  @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3ComponentOverrideApi public final class BasicAlertDialogOverrideScope {
+    method public kotlin.jvm.functions.Function0<kotlin.Unit> getContent();
+    method public androidx.compose.ui.Modifier getModifier();
+    method public kotlin.jvm.functions.Function0<kotlin.Unit> getOnDismissRequest();
+    method public androidx.compose.ui.window.DialogProperties getProperties();
+    property public kotlin.jvm.functions.Function0<kotlin.Unit> content;
+    property public androidx.compose.ui.Modifier modifier;
+    property public kotlin.jvm.functions.Function0<kotlin.Unit> onDismissRequest;
+    property public androidx.compose.ui.window.DialogProperties properties;
+  }
+
   public final class BottomAppBarDefaults {
     method @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3Api @androidx.compose.runtime.Composable public androidx.compose.material3.BottomAppBarScrollBehavior exitAlwaysScrollBehavior(optional androidx.compose.material3.BottomAppBarState state, optional kotlin.jvm.functions.Function0<java.lang.Boolean> canScroll, optional androidx.compose.animation.core.AnimationSpec<java.lang.Float>? snapAnimationSpec, optional androidx.compose.animation.core.DecayAnimationSpec<java.lang.Float>? flingAnimationSpec);
     method @androidx.compose.runtime.Composable public long getBottomAppBarFabColor();
@@ -721,6 +738,11 @@
     property public abstract kotlin.ranges.IntRange yearRange;
   }
 
+  @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3ComponentOverrideApi public final class DefaultBasicAlertDialogOverride implements androidx.compose.material3.BasicAlertDialogOverride {
+    method @androidx.compose.runtime.Composable public void BasicAlertDialog(androidx.compose.material3.BasicAlertDialogOverrideScope);
+    field public static final androidx.compose.material3.DefaultBasicAlertDialogOverride INSTANCE;
+  }
+
   @SuppressCompatibility @androidx.compose.material3.ExperimentalMaterial3ComponentOverrideApi public final class DefaultNavigationBarComponentOverride implements androidx.compose.material3.NavigationBarComponentOverride {
     method @androidx.compose.runtime.Composable public void NavigationBar(androidx.compose.material3.NavigationBarComponentOverrideContext);
     field public static final androidx.compose.material3.DefaultNavigationBarComponentOverride INSTANCE;
diff --git a/compose/material3/material3/bcv/native/current.txt b/compose/material3/material3/bcv/native/current.txt
index 1072e3c..2faced8 100644
--- a/compose/material3/material3/bcv/native/current.txt
+++ b/compose/material3/material3/bcv/native/current.txt
@@ -109,6 +109,10 @@
     abstract suspend fun snapTo(kotlin/Float) // androidx.compose.material3.pulltorefresh/PullToRefreshState.snapTo|snapTo(kotlin.Float){}[0]
 }
 
+abstract interface androidx.compose.material3/BasicAlertDialogOverride { // androidx.compose.material3/BasicAlertDialogOverride|null[0]
+    abstract fun (androidx.compose.material3/BasicAlertDialogOverrideScope).BasicAlertDialog() // androidx.compose.material3/BasicAlertDialogOverride.BasicAlertDialog|BasicAlertDialog@androidx.compose.material3.BasicAlertDialogOverrideScope(){}[0]
+}
+
 abstract interface androidx.compose.material3/BottomAppBarScrollBehavior { // androidx.compose.material3/BottomAppBarScrollBehavior|null[0]
     abstract val flingAnimationSpec // androidx.compose.material3/BottomAppBarScrollBehavior.flingAnimationSpec|{}flingAnimationSpec[0]
         abstract fun <get-flingAnimationSpec>(): androidx.compose.animation.core/DecayAnimationSpec<kotlin/Float>? // androidx.compose.material3/BottomAppBarScrollBehavior.flingAnimationSpec.<get-flingAnimationSpec>|<get-flingAnimationSpec>(){}[0]
@@ -432,6 +436,17 @@
     }
 }
 
+final class androidx.compose.material3/BasicAlertDialogOverrideScope { // androidx.compose.material3/BasicAlertDialogOverrideScope|null[0]
+    final val content // androidx.compose.material3/BasicAlertDialogOverrideScope.content|{}content[0]
+        final fun <get-content>(): kotlin/Function0<kotlin/Unit> // androidx.compose.material3/BasicAlertDialogOverrideScope.content.<get-content>|<get-content>(){}[0]
+    final val modifier // androidx.compose.material3/BasicAlertDialogOverrideScope.modifier|{}modifier[0]
+        final fun <get-modifier>(): androidx.compose.ui/Modifier // androidx.compose.material3/BasicAlertDialogOverrideScope.modifier.<get-modifier>|<get-modifier>(){}[0]
+    final val onDismissRequest // androidx.compose.material3/BasicAlertDialogOverrideScope.onDismissRequest|{}onDismissRequest[0]
+        final fun <get-onDismissRequest>(): kotlin/Function0<kotlin/Unit> // androidx.compose.material3/BasicAlertDialogOverrideScope.onDismissRequest.<get-onDismissRequest>|<get-onDismissRequest>(){}[0]
+    final val properties // androidx.compose.material3/BasicAlertDialogOverrideScope.properties|{}properties[0]
+        final fun <get-properties>(): androidx.compose.ui.window/DialogProperties // androidx.compose.material3/BasicAlertDialogOverrideScope.properties.<get-properties>|<get-properties>(){}[0]
+}
+
 final class androidx.compose.material3/BottomSheetScaffoldState { // androidx.compose.material3/BottomSheetScaffoldState|null[0]
     constructor <init>(androidx.compose.material3/SheetState, androidx.compose.material3/SnackbarHostState) // androidx.compose.material3/BottomSheetScaffoldState.<init>|<init>(androidx.compose.material3.SheetState;androidx.compose.material3.SnackbarHostState){}[0]
 
@@ -2289,6 +2304,10 @@
     final fun DateRangePickerTitle(androidx.compose.material3/DisplayMode, androidx.compose.ui/Modifier = ..., androidx.compose.ui.graphics/Color = ...) // androidx.compose.material3/DateRangePickerDefaults.DateRangePickerTitle|DateRangePickerTitle(androidx.compose.material3.DisplayMode;androidx.compose.ui.Modifier;androidx.compose.ui.graphics.Color){}[0]
 }
 
+final object androidx.compose.material3/DefaultBasicAlertDialogOverride : androidx.compose.material3/BasicAlertDialogOverride { // androidx.compose.material3/DefaultBasicAlertDialogOverride|null[0]
+    final fun (androidx.compose.material3/BasicAlertDialogOverrideScope).BasicAlertDialog() // androidx.compose.material3/DefaultBasicAlertDialogOverride.BasicAlertDialog|BasicAlertDialog@androidx.compose.material3.BasicAlertDialogOverrideScope(){}[0]
+}
+
 final object androidx.compose.material3/DefaultNavigationBarComponentOverride : androidx.compose.material3/NavigationBarComponentOverride { // androidx.compose.material3/DefaultNavigationBarComponentOverride|null[0]
     final fun (androidx.compose.material3/NavigationBarComponentOverrideContext).NavigationBar() // androidx.compose.material3/DefaultNavigationBarComponentOverride.NavigationBar|NavigationBar@androidx.compose.material3.NavigationBarComponentOverrideContext(){}[0]
 }
@@ -3197,6 +3216,8 @@
 
 final val androidx.compose.material3/LocalAbsoluteTonalElevation // androidx.compose.material3/LocalAbsoluteTonalElevation|{}LocalAbsoluteTonalElevation[0]
     final fun <get-LocalAbsoluteTonalElevation>(): androidx.compose.runtime/ProvidableCompositionLocal<androidx.compose.ui.unit/Dp> // androidx.compose.material3/LocalAbsoluteTonalElevation.<get-LocalAbsoluteTonalElevation>|<get-LocalAbsoluteTonalElevation>(){}[0]
+final val androidx.compose.material3/LocalBasicAlertDialogOverride // androidx.compose.material3/LocalBasicAlertDialogOverride|{}LocalBasicAlertDialogOverride[0]
+    final fun <get-LocalBasicAlertDialogOverride>(): androidx.compose.runtime/ProvidableCompositionLocal<androidx.compose.material3/BasicAlertDialogOverride> // androidx.compose.material3/LocalBasicAlertDialogOverride.<get-LocalBasicAlertDialogOverride>|<get-LocalBasicAlertDialogOverride>(){}[0]
 final val androidx.compose.material3/LocalContentColor // androidx.compose.material3/LocalContentColor|{}LocalContentColor[0]
     final fun <get-LocalContentColor>(): androidx.compose.runtime/ProvidableCompositionLocal<androidx.compose.ui.graphics/Color> // androidx.compose.material3/LocalContentColor.<get-LocalContentColor>|<get-LocalContentColor>(){}[0]
 final val androidx.compose.material3/LocalMinimumInteractiveComponentEnforcement // androidx.compose.material3/LocalMinimumInteractiveComponentEnforcement|{}LocalMinimumInteractiveComponentEnforcement[0]
diff --git a/compose/material3/material3/src/commonMain/kotlin/androidx/compose/material3/AlertDialog.kt b/compose/material3/material3/src/commonMain/kotlin/androidx/compose/material3/AlertDialog.kt
index 091e453..450a64d 100644
--- a/compose/material3/material3/src/commonMain/kotlin/androidx/compose/material3/AlertDialog.kt
+++ b/compose/material3/material3/src/commonMain/kotlin/androidx/compose/material3/AlertDialog.kt
@@ -28,6 +28,8 @@
 import androidx.compose.material3.tokens.DialogTokens
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.ProvidableCompositionLocal
+import androidx.compose.runtime.compositionLocalOf
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
@@ -134,6 +136,7 @@
  * @param properties typically platform specific properties to further configure the dialog.
  * @param content the content of the dialog
  */
+@OptIn(ExperimentalMaterial3ComponentOverrideApi::class)
 @ExperimentalMaterial3Api
 @Composable
 fun BasicAlertDialog(
@@ -142,19 +145,41 @@
     properties: DialogProperties = DialogProperties(),
     content: @Composable () -> Unit
 ) {
-    Dialog(
-        onDismissRequest = onDismissRequest,
-        properties = properties,
-    ) {
-        val dialogPaneDescription = getString(Strings.Dialog)
-        Box(
-            modifier =
-                modifier
-                    .sizeIn(minWidth = DialogMinWidth, maxWidth = DialogMaxWidth)
-                    .then(Modifier.semantics { paneTitle = dialogPaneDescription }),
-            propagateMinConstraints = true
+    with(LocalBasicAlertDialogOverride.current) {
+        BasicAlertDialogOverrideScope(
+                onDismissRequest = onDismissRequest,
+                modifier = modifier,
+                properties = properties,
+                content = content
+            )
+            .BasicAlertDialog()
+    }
+}
+
+/**
+ * This override provides the default behavior of the [BasicAlertDialog] component.
+ *
+ * [BasicAlertDialogOverride] used when no override is specified.
+ */
+@OptIn(ExperimentalMaterial3Api::class)
+@ExperimentalMaterial3ComponentOverrideApi
+object DefaultBasicAlertDialogOverride : BasicAlertDialogOverride {
+    @Composable
+    override fun BasicAlertDialogOverrideScope.BasicAlertDialog() {
+        Dialog(
+            onDismissRequest = onDismissRequest,
+            properties = properties,
         ) {
-            content()
+            val dialogPaneDescription = getString(Strings.Dialog)
+            Box(
+                modifier =
+                    modifier
+                        .sizeIn(minWidth = DialogMinWidth, maxWidth = DialogMaxWidth)
+                        .then(Modifier.semantics { paneTitle = dialogPaneDescription }),
+                propagateMinConstraints = true
+            ) {
+                content()
+            }
         }
     }
 }
@@ -241,7 +266,7 @@
     textContentColor: Color,
     tonalElevation: Dp,
     properties: DialogProperties
-): Unit =
+) {
     BasicAlertDialog(
         onDismissRequest = onDismissRequest,
         modifier = modifier,
@@ -264,15 +289,16 @@
             containerColor = containerColor,
             tonalElevation = tonalElevation,
             // Note that a button content color is provided here from the dialog's token, but in
-            // most cases, TextButtons should be used for dismiss and confirm buttons.
-            // TextButtons will not consume this provided content color value, and will used their
-            // own defined or default colors.
+            // most cases, TextButtons should be used for dismiss and confirm buttons. TextButtons
+            // will not consume this provided content color value, and will used their own defined
+            // or default colors.
             buttonContentColor = DialogTokens.ActionLabelTextColor.value,
             iconContentColor = iconContentColor,
             titleContentColor = titleContentColor,
             textContentColor = textContentColor,
         )
     }
+}
 
 @Composable
 internal fun AlertDialogContent(
@@ -458,3 +484,42 @@
 private val IconPadding = PaddingValues(bottom = 16.dp)
 private val TitlePadding = PaddingValues(bottom = 16.dp)
 private val TextPadding = PaddingValues(bottom = 24.dp)
+
+/**
+ * Interface that allows libraries to override the behavior of the [BasicAlertDialog] component.
+ *
+ * To override this component, implement the member function of this interface, then provide the
+ * implementation to [LocalBasicAlertDialogOverride] in the Compose hierarchy.
+ */
+@ExperimentalMaterial3ComponentOverrideApi
+interface BasicAlertDialogOverride {
+    /** Behavior function that is called by the [BasicAlertDialog] component. */
+    @Composable fun BasicAlertDialogOverrideScope.BasicAlertDialog()
+}
+
+/**
+ * Parameters available to [BasicAlertDialog].
+ *
+ * @param onDismissRequest called when the user tries to dismiss the Dialog by clicking outside or
+ *   pressing the back button. This is not called when the dismiss button is clicked.
+ * @param modifier the [Modifier] to be applied to this dialog's content.
+ * @param properties typically platform specific properties to further configure the dialog.
+ * @param content the content of the dialog
+ */
+@ExperimentalMaterial3ComponentOverrideApi
+class BasicAlertDialogOverrideScope
+internal constructor(
+    val onDismissRequest: () -> Unit,
+    val modifier: Modifier = Modifier,
+    val properties: DialogProperties = DialogProperties(),
+    val content: @Composable () -> Unit
+)
+
+/** CompositionLocal containing the currently-selected [BasicAlertDialogOverride]. */
+@Suppress("OPT_IN_MARKER_ON_WRONG_TARGET")
+@get:ExperimentalMaterial3ComponentOverrideApi
+@ExperimentalMaterial3ComponentOverrideApi
+val LocalBasicAlertDialogOverride: ProvidableCompositionLocal<BasicAlertDialogOverride> =
+    compositionLocalOf {
+        DefaultBasicAlertDialogOverride
+    }
diff --git a/compose/material3/material3/src/commonMain/kotlin/androidx/compose/material3/ButtonGroup.kt b/compose/material3/material3/src/commonMain/kotlin/androidx/compose/material3/ButtonGroup.kt
index 50375e5..a52553e 100644
--- a/compose/material3/material3/src/commonMain/kotlin/androidx/compose/material3/ButtonGroup.kt
+++ b/compose/material3/material3/src/commonMain/kotlin/androidx/compose/material3/ButtonGroup.kt
@@ -25,6 +25,7 @@
 import androidx.compose.foundation.layout.Arrangement
 import androidx.compose.foundation.shape.RoundedCornerShape
 import androidx.compose.material3.tokens.ButtonGroupSmallTokens
+import androidx.compose.material3.tokens.ConnectedButtonGroupSmallTokens
 import androidx.compose.material3.tokens.MotionSchemeKeyTokens
 import androidx.compose.material3.tokens.ShapeTokens
 import androidx.compose.runtime.Composable
@@ -49,7 +50,6 @@
 import androidx.compose.ui.unit.Constraints
 import androidx.compose.ui.unit.Density
 import androidx.compose.ui.unit.Dp
-import androidx.compose.ui.unit.dp
 import androidx.compose.ui.util.fastForEach
 import androidx.compose.ui.util.fastForEachIndexed
 import androidx.compose.ui.util.fastMapIndexed
@@ -192,55 +192,50 @@
         Arrangement.spacedBy(ButtonGroupSmallTokens.BetweenSpace)
 
     /** The default spacing used between children for connected button group */
-    // TODO replace with token value
-    val ConnectedSpaceBetween: Dp = 2.dp
+    val ConnectedSpaceBetween: Dp = ConnectedButtonGroupSmallTokens.BetweenSpace
 
     /** Default shape for the leading button in a connected button group */
     val connectedLeadingButtonShape: Shape
         @Composable
         get() =
-            // TODO replace with token value
             RoundedCornerShape(
                 topStart = ShapeDefaults.CornerFull,
                 bottomStart = ShapeDefaults.CornerFull,
-                topEnd = ShapeDefaults.CornerSmall,
-                bottomEnd = ShapeDefaults.CornerSmall
+                topEnd = ConnectedButtonGroupSmallTokens.InnerCornerCornerSize,
+                bottomEnd = ConnectedButtonGroupSmallTokens.InnerCornerCornerSize
             )
 
     /** Default shape for the pressed state for the leading button in a connected button group. */
     val connectedLeadingButtonPressShape: Shape
         @Composable
         get() =
-            // TODO replace with token value
             RoundedCornerShape(
                 topStart = ShapeDefaults.CornerFull,
                 bottomStart = ShapeDefaults.CornerFull,
-                topEnd = ShapeDefaults.CornerExtraSmall,
-                bottomEnd = ShapeDefaults.CornerExtraSmall
+                topEnd = ConnectedButtonGroupSmallTokens.PressedInnerCornerCornerSize,
+                bottomEnd = ConnectedButtonGroupSmallTokens.PressedInnerCornerCornerSize
             )
 
     /** Default shape for the trailing button in a connected button group */
     val connectedTrailingButtonShape: Shape
         @Composable
         get() =
-            // TODO replace with token value
             RoundedCornerShape(
                 topEnd = ShapeDefaults.CornerFull,
                 bottomEnd = ShapeDefaults.CornerFull,
-                topStart = ShapeDefaults.CornerSmall,
-                bottomStart = ShapeDefaults.CornerSmall
+                topStart = ConnectedButtonGroupSmallTokens.InnerCornerCornerSize,
+                bottomStart = ConnectedButtonGroupSmallTokens.InnerCornerCornerSize
             )
 
     /** Default shape for the pressed state for the trailing button in a connected button group. */
     val connectedTrailingButtonPressShape: Shape
         @Composable
         get() =
-            // TODO replace with token value
             RoundedCornerShape(
                 topEnd = ShapeDefaults.CornerFull,
                 bottomEnd = ShapeDefaults.CornerFull,
-                topStart = ShapeDefaults.CornerExtraSmall,
-                bottomStart = ShapeDefaults.CornerExtraSmall
+                topStart = ConnectedButtonGroupSmallTokens.PressedInnerCornerCornerSize,
+                bottomStart = ConnectedButtonGroupSmallTokens.PressedInnerCornerCornerSize
             )
 
     /** Default shape for the checked state for the buttons in a connected button group */
@@ -249,9 +244,7 @@
     /** Default shape for the pressed state for the middle buttons in a connected button group. */
     val connectedMiddleButtonPressShape: Shape
         @Composable
-        get() =
-            // TODO replace with token value
-            RoundedCornerShape(ShapeDefaults.CornerExtraSmall)
+        get() = RoundedCornerShape(ConnectedButtonGroupSmallTokens.PressedInnerCornerCornerSize)
 
     /** Defaults button shapes for the start button in a [ConnectedButtonGroup] */
     @Composable
diff --git a/compose/material3/material3/src/commonMain/kotlin/androidx/compose/material3/tokens/ConnectedButtonGroupSmallTokens.kt b/compose/material3/material3/src/commonMain/kotlin/androidx/compose/material3/tokens/ConnectedButtonGroupSmallTokens.kt
new file mode 100644
index 0000000..cf528a8
--- /dev/null
+++ b/compose/material3/material3/src/commonMain/kotlin/androidx/compose/material3/tokens/ConnectedButtonGroupSmallTokens.kt
@@ -0,0 +1,31 @@
+/*
+ * Copyright 2025 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.
+ */
+
+// VERSION: 14_1_0
+// GENERATED CODE - DO NOT MODIFY BY HAND
+
+package androidx.compose.material3.tokens
+
+import androidx.compose.ui.unit.dp
+
+internal object ConnectedButtonGroupSmallTokens {
+    val BetweenSpace = 2.0.dp
+    val ContainerHeight = 40.0.dp
+    val ContainerShape = ShapeKeyTokens.CornerFull
+    val InnerCornerCornerSize = ShapeTokens.CornerValueSmall
+    val PressedInnerCornerCornerSize = ShapeTokens.CornerValueExtraSmall
+    val SelectedInnerCornerCornerSizePercent = 50.0f
+}
diff --git a/libraryversions.toml b/libraryversions.toml
index 152b2a3..b4c390b 100644
--- a/libraryversions.toml
+++ b/libraryversions.toml
@@ -20,7 +20,7 @@
 CAMERA_VIEWFINDER = "1.4.0-alpha13"
 CARDVIEW = "1.1.0-alpha01"
 CAR_APP = "1.8.0-alpha01"
-COLLECTION = "1.5.0-rc01"
+COLLECTION = "1.5.0"
 COMPOSE = "1.9.0-alpha01"
 COMPOSE_MATERIAL3 = "1.4.0-alpha09"
 COMPOSE_MATERIAL3_ADAPTIVE = "1.2.0-alpha01"
@@ -107,7 +107,7 @@
 NAVIGATION3 = "0.1.0-dev01"
 PAGING = "3.4.0-alpha01"
 PALETTE = "1.1.0-alpha01"
-PDF = "1.0.0-alpha07"
+PDF = "1.0.0-alpha08"
 PERCENTLAYOUT = "1.1.0-alpha01"
 PERFORMANCE = "1.0.0-alpha01"
 PREFERENCE = "1.3.0-alpha01"
diff --git a/lifecycle/lifecycle-viewmodel-navigation3/src/androidInstrumentedTest/kotlin/androidx/lifecycle/viewmodel/navigation3/ViewModelStoreNavLocalProviderTest.kt b/lifecycle/lifecycle-viewmodel-navigation3/src/androidInstrumentedTest/kotlin/androidx/lifecycle/viewmodel/navigation3/ViewModelStoreNavLocalProviderTest.kt
index 0476584..e88cccc 100644
--- a/lifecycle/lifecycle-viewmodel-navigation3/src/androidInstrumentedTest/kotlin/androidx/lifecycle/viewmodel/navigation3/ViewModelStoreNavLocalProviderTest.kt
+++ b/lifecycle/lifecycle-viewmodel-navigation3/src/androidInstrumentedTest/kotlin/androidx/lifecycle/viewmodel/navigation3/ViewModelStoreNavLocalProviderTest.kt
@@ -58,8 +58,8 @@
                 viewModel2.myArg = entry2Arg
             }
         composeTestRule.setContent {
-            savedStateWrapper.ProvideToBackStack(backStack = listOf(entry1, entry2)) {
-                viewModelWrapper.ProvideToBackStack(backStack = listOf(entry1, entry2)) {
+            savedStateWrapper.ProvideToBackStack(backStack = listOf(entry1.key, entry2.key)) {
+                viewModelWrapper.ProvideToBackStack(backStack = listOf(entry1.key, entry2.key)) {
                     savedStateWrapper.ProvideToEntry(
                         NavEntry(entry1.key) { viewModelWrapper.ProvideToEntry(entry1) }
                     )
@@ -92,7 +92,7 @@
             }
         try {
             composeTestRule.setContent {
-                viewModelWrapper.ProvideToBackStack(backStack = listOf(entry1)) {
+                viewModelWrapper.ProvideToBackStack(backStack = listOf(entry1.key)) {
                     viewModelWrapper.ProvideToEntry(entry1)
                 }
             }
diff --git a/lifecycle/lifecycle-viewmodel-navigation3/src/androidMain/kotlin/androidx/lifecycle/viewmodel/navigation3/ViewModelStoreNavLocalProvider.android.kt b/lifecycle/lifecycle-viewmodel-navigation3/src/androidMain/kotlin/androidx/lifecycle/viewmodel/navigation3/ViewModelStoreNavLocalProvider.android.kt
index 5c4b682..5bad76e 100644
--- a/lifecycle/lifecycle-viewmodel-navigation3/src/androidMain/kotlin/androidx/lifecycle/viewmodel/navigation3/ViewModelStoreNavLocalProvider.android.kt
+++ b/lifecycle/lifecycle-viewmodel-navigation3/src/androidMain/kotlin/androidx/lifecycle/viewmodel/navigation3/ViewModelStoreNavLocalProvider.android.kt
@@ -17,11 +17,11 @@
 package androidx.lifecycle.viewmodel.navigation3
 
 import androidx.activity.compose.LocalActivity
-import androidx.collection.MutableObjectIntMap
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
 import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.remember
+import androidx.compose.runtime.saveable.rememberSaveable
 import androidx.compose.runtime.staticCompositionLocalOf
 import androidx.lifecycle.HasDefaultViewModelProviderFactory
 import androidx.lifecycle.Lifecycle
@@ -57,33 +57,42 @@
         entryViewModelStoreProvider.ownerInBackStack.clear()
         entryViewModelStoreProvider.ownerInBackStack.addAll(backStack)
         val localInfo = remember { ViewModelStoreNavLocalInfo() }
-        DisposableEffect(key1 = backStack) {
-            localInfo.refCount.clear()
-            onDispose {}
-        }
+        DisposableEffect(key1 = backStack) { onDispose { localInfo.refCount.clear() } }
 
         val activity = LocalActivity.current
-        backStack.forEach { key ->
+        backStack.forEachIndexed { index, key ->
+            // We update here as part of composition to ensure the value is available to
+            // ProvideToEntry
+            localInfo.refCount.getOrPut(key) { LinkedHashSet<Int>() }.add(getIdForKey(key, index))
             DisposableEffect(key1 = key) {
-                localInfo.refCount[key] = localInfo.refCount.getOrDefault(key, 0).plus(1)
+                localInfo.refCount
+                    .getOrPut(key) { LinkedHashSet<Int>() }
+                    .add(getIdForKey(key, index))
                 onDispose {
-                    localInfo.refCount[key] =
-                        localInfo.refCount
-                            .getOrElse(key) {
-                                error(
-                                    "Attempting to incorrectly dispose of backstack state in " +
-                                        "ViewModelStoreNavLocalProvider"
-                                )
+                    // If the backStack count is less than the refCount for the key, remove the
+                    // state since that means we removed a key from the backstack, and set the
+                    // refCount to the backstack count.
+                    val backstackCount = backStack.count { it == key }
+                    val lastKeyCount = localInfo.refCount[key]?.size ?: 0
+                    if (backstackCount < lastKeyCount) {
+                        // The set of the ids associated with this key
+                        @Suppress("PrimitiveInCollection") // The order of the element matters
+                        val idsSet = localInfo.refCount[key]!!
+                        val id = idsSet.last()
+                        idsSet.remove(id)
+                        if (!localInfo.idsInComposition.contains(id)) {
+                            if (activity?.isChangingConfigurations != true) {
+                                entryViewModelStoreProvider
+                                    .removeViewModelStoreOwnerForKey(id)
+                                    ?.clear()
                             }
-                            .minus(1)
-                    if (localInfo.refCount[key] <= 0) {
-                        // This ensures we always keep viewModels on config changes.
-                        if (activity?.isChangingConfigurations != true) {
-                            entryViewModelStoreProvider
-                                .removeViewModelStoreOwnerForKey(key)
-                                ?.clear()
                         }
                     }
+
+                    // If the refCount is 0, remove the key from the refCount.
+                    if (localInfo.refCount[key]?.isEmpty() == true) {
+                        localInfo.refCount.remove(key)
+                    }
                 }
             }
         }
@@ -97,34 +106,48 @@
     override fun <T : Any> ProvideToEntry(entry: NavEntry<T>) {
         val key = entry.key
         val entryViewModelStoreProvider = viewModel { EntryViewModel() }
-        val viewModelStore = entryViewModelStoreProvider.viewModelStoreForKey(key)
 
         val activity = LocalActivity.current
         val localInfo = LocalViewModelStoreNavLocalInfo.current
+        // Tracks whether the key is changed
+        var keyChanged = false
+        var id: Int =
+            rememberSaveable(key) {
+                keyChanged = true
+                localInfo.refCount[key]!!.last()
+            }
+        id =
+            rememberSaveable(localInfo.refCount[key]?.size) {
+                // if the key changed, use the current id
+                // If the key was not changed, and the current id is not in composition or on the
+                // back stack then update the id with the last item from the backstack with the
+                // associated key. This ensures that we can handle duplicates, both consecutive and
+                // non-consecutive
+                if (
+                    !keyChanged &&
+                        (!localInfo.idsInComposition.contains(id) ||
+                            localInfo.refCount[key]?.contains(id) == true)
+                ) {
+                    localInfo.refCount[key]!!.last()
+                } else {
+                    id
+                }
+            }
+        keyChanged = false
+
+        val viewModelStore = entryViewModelStoreProvider.viewModelStoreForKey(id)
 
         DisposableEffect(key1 = key) {
-            localInfo.refCount[key] = localInfo.refCount.getOrDefault(key, 0).plus(1)
+            localInfo.idsInComposition.add(id)
             onDispose {
-                // We need to check to make sure that the refcount has been cleared here because
-                // when we are using animations, if the entire back stack is changed, we will
-                // execute the onDispose above that clears all of the counts before we finish the
-                // transition and run this onDispose so our count will already be gone and we
-                // should just remove the state.
-                if (!localInfo.refCount.contains(key) || localInfo.refCount[key] == 0) {
-                    // This ensures we always keep viewModels on config changes.
+                if (localInfo.idsInComposition.remove(id) && !localInfo.refCount.contains(key)) {
                     if (activity?.isChangingConfigurations != true) {
-                        entryViewModelStoreProvider.removeViewModelStoreOwnerForKey(key)?.clear()
+                        entryViewModelStoreProvider.removeViewModelStoreOwnerForKey(id)?.clear()
                     }
-                } else {
-                    localInfo.refCount[key] =
-                        localInfo.refCount
-                            .getOrElse(key) {
-                                error(
-                                    "Attempting to incorrectly dispose of state associated with " +
-                                        "key $key in ViewModelStoreNavLocalProvider."
-                                )
-                            }
-                            .minus(1)
+                    // If the refCount is 0, remove the key from the refCount.
+                    if (localInfo.refCount[key]?.isEmpty() == true) {
+                        localInfo.refCount.remove(key)
+                    }
                 }
             }
         }
@@ -187,5 +210,9 @@
     }
 
 internal class ViewModelStoreNavLocalInfo {
-    internal val refCount: MutableObjectIntMap<Any> = MutableObjectIntMap()
+    internal val refCount: MutableMap<Any, LinkedHashSet<Int>> = mutableMapOf()
+    @Suppress("PrimitiveInCollection") // The order of the element matters
+    internal val idsInComposition: LinkedHashSet<Int> = LinkedHashSet<Int>()
 }
+
+internal fun getIdForKey(key: Any, count: Int): Int = 31 * key.hashCode() + count
diff --git a/navigation3/navigation3/src/androidInstrumentedTest/kotlin/androidx/navigation3/AnimatedTest.kt b/navigation3/navigation3/src/androidInstrumentedTest/kotlin/androidx/navigation3/AnimatedTest.kt
index e227982..c963e01 100644
--- a/navigation3/navigation3/src/androidInstrumentedTest/kotlin/androidx/navigation3/AnimatedTest.kt
+++ b/navigation3/navigation3/src/androidInstrumentedTest/kotlin/androidx/navigation3/AnimatedTest.kt
@@ -23,10 +23,17 @@
 import androidx.compose.animation.core.tween
 import androidx.compose.animation.fadeIn
 import androidx.compose.animation.fadeOut
+import androidx.compose.foundation.background
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.material3.Text
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.compositionLocalOf
 import androidx.compose.runtime.mutableStateListOf
 import androidx.compose.runtime.remember
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.test.assertCountEquals
 import androidx.compose.ui.test.assertIsDisplayed
 import androidx.compose.ui.test.isDisplayed
@@ -721,6 +728,105 @@
     }
 
     @Test
+    fun testPoppedEntryIsAnimated() {
+        lateinit var backstack: MutableList<Any>
+        composeTestRule.setContent {
+            backstack = remember { mutableStateListOf(first, second) }
+            SinglePaneNavDisplay(backstack) {
+                when (it) {
+                    first ->
+                        NavEntry(
+                            first,
+                        ) {
+                            Text(first)
+                        }
+                    second ->
+                        NavEntry(
+                            second,
+                        ) {
+                            Box(Modifier.fillMaxSize().background(Color.Red)) { Text(second) }
+                        }
+                    else -> error("Invalid key passed")
+                }
+            }
+        }
+
+        composeTestRule.waitForIdle()
+        composeTestRule.onNodeWithText(second).assertIsDisplayed()
+        assertThat(backstack).containsExactly(first, second)
+
+        composeTestRule.mainClock.autoAdvance = false
+        composeTestRule.runOnIdle { backstack.removeAt(1) }
+
+        // advance by a duration that is much shorter than the default duration
+        // to ensure that the custom animation is used and has completed after this
+        composeTestRule.mainClock.advanceTimeByFrame()
+        composeTestRule.mainClock.advanceTimeByFrame()
+
+        composeTestRule.onNodeWithText(second).assertIsDisplayed()
+    }
+
+    @Test
+    fun testPoppedEntryIsWrapped() {
+        lateinit var backstack: MutableList<Any>
+        val LocalHasProvidedToEntry = compositionLocalOf { false }
+        val provider =
+            object : NavLocalProvider {
+                @Composable
+                override fun ProvideToBackStack(
+                    backStack: List<Any>,
+                    content: @Composable () -> Unit
+                ) {
+                    CompositionLocalProvider(LocalHasProvidedToEntry provides false) {
+                        content.invoke()
+                    }
+                }
+
+                @Composable
+                override fun <T : Any> ProvideToEntry(entry: NavEntry<T>) {
+                    CompositionLocalProvider(LocalHasProvidedToEntry provides true) {
+                        entry.content.invoke(entry.key)
+                    }
+                }
+            }
+        var secondEntryIsWrapped = false
+        composeTestRule.setContent {
+            backstack = remember { mutableStateListOf(first, second) }
+            SinglePaneNavDisplay(backstack, localProviders = listOf(provider)) {
+                when (it) {
+                    first ->
+                        NavEntry(
+                            first,
+                        ) {
+                            Text(first)
+                        }
+                    second ->
+                        NavEntry(
+                            second,
+                        ) {
+                            secondEntryIsWrapped = LocalHasProvidedToEntry.current
+                            Box(Modifier.fillMaxSize().background(Color.Red)) { Text(second) }
+                        }
+                    else -> error("Invalid key passed")
+                }
+            }
+        }
+
+        composeTestRule.waitForIdle()
+        composeTestRule.onNodeWithText(second).assertIsDisplayed()
+        assertThat(backstack).containsExactly(first, second)
+
+        composeTestRule.mainClock.autoAdvance = false
+        composeTestRule.runOnIdle { backstack.removeAt(1) }
+
+        // advance by a duration that is much shorter than the default duration
+        // to ensure that the custom animation is used and has completed after this
+        composeTestRule.mainClock.advanceTimeByFrame()
+        composeTestRule.mainClock.advanceTimeByFrame()
+        assertTrue(secondEntryIsWrapped)
+    }
+
+    @Test
     fun testDuplicateLastEntry() {
         lateinit var backStack: MutableList<Any>
         val testDuration = DEFAULT_TRANSITION_DURATION_MILLISECOND / 5
diff --git a/navigation3/navigation3/src/androidMain/kotlin/androidx/navigation3/SinglePaneNavDisplay.android.kt b/navigation3/navigation3/src/androidMain/kotlin/androidx/navigation3/SinglePaneNavDisplay.android.kt
index 6a40e5e..1878e88 100644
--- a/navigation3/navigation3/src/androidMain/kotlin/androidx/navigation3/SinglePaneNavDisplay.android.kt
+++ b/navigation3/navigation3/src/androidMain/kotlin/androidx/navigation3/SinglePaneNavDisplay.android.kt
@@ -232,12 +232,13 @@
             contentAlignment = contentAlignment,
             contentKey = { it.last() }
         ) { innerStack ->
-            val lastKey = innerStack.last()
-            val lastEntry = entries.findLast { entry -> entry.key == lastKey }
-            lastEntry?.let { entry ->
-                CompositionLocalProvider(LocalNavAnimatedContentScope provides this) {
-                    entry.content.invoke(lastKey)
+            // popped entries will be remembered and retrieved so it can animate out properly
+            val currEntry =
+                remember((innerStack.last())) {
+                    entries.findLast { entry -> entry.key == innerStack.last() }
                 }
+            CompositionLocalProvider(LocalNavAnimatedContentScope provides this) {
+                currEntry!!.content.invoke(currEntry.key)
             }
         }
     }
diff --git a/navigation3/navigation3/src/commonMain/kotlin/androidx/navigation3/SaveableStateNavLocalProvider.kt b/navigation3/navigation3/src/commonMain/kotlin/androidx/navigation3/SaveableStateNavLocalProvider.kt
index 1ebca75..3fa9353 100644
--- a/navigation3/navigation3/src/commonMain/kotlin/androidx/navigation3/SaveableStateNavLocalProvider.kt
+++ b/navigation3/navigation3/src/commonMain/kotlin/androidx/navigation3/SaveableStateNavLocalProvider.kt
@@ -16,7 +16,6 @@
 
 package androidx.navigation3
 
-import androidx.collection.MutableObjectIntMap
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
 import androidx.compose.runtime.DisposableEffect
@@ -25,6 +24,7 @@
 import androidx.compose.runtime.saveable.rememberSaveable
 import androidx.compose.runtime.saveable.rememberSaveableStateHolder
 import androidx.compose.runtime.staticCompositionLocalOf
+import kotlin.collections.LinkedHashSet
 
 /**
  * Wraps the content of a [NavEntry] with a [SaveableStateHolder.SaveableStateProvider] to ensure
@@ -38,33 +38,37 @@
     @Composable
     override fun ProvideToBackStack(backStack: List<Any>, content: @Composable () -> Unit) {
         val localInfo = remember { SaveableStateNavLocalInfo() }
-        DisposableEffect(key1 = backStack) {
-            localInfo.refCount.clear()
-            onDispose {}
-        }
+        DisposableEffect(key1 = backStack) { onDispose { localInfo.refCount.clear() } }
 
         localInfo.savedStateHolder = rememberSaveableStateHolder()
-        backStack.forEach { key ->
+        backStack.forEachIndexed { index, key ->
             // We update here as part of composition to ensure the value is available to
             // ProvideToEntry
-            localInfo.refCount[key] = backStack.count { it == key }
+            localInfo.refCount.getOrPut(key) { LinkedHashSet<Int>() }.add(getIdForKey(key, index))
             DisposableEffect(key1 = key) {
                 // We update here at the end of composition in case the backstack changed and
                 // everything was cleared.
-                localInfo.refCount[key] = backStack.count { it == key }
+                localInfo.refCount
+                    .getOrPut(key) { LinkedHashSet<Int>() }
+                    .add(getIdForKey(key, index))
                 onDispose {
                     // If the backStack count is less than the refCount for the key, remove the
                     // state since that means we removed a key from the backstack, and set the
                     // refCount to the backstack count.
                     val backstackCount = backStack.count { it == key }
-                    if (backstackCount < localInfo.refCount[key]) {
-                        localInfo.savedStateHolder!!.removeState(
-                            getIdForKey(key, localInfo.refCount[key])
-                        )
-                        localInfo.refCount[key] = backstackCount
+                    val lastKeyCount = localInfo.refCount[key]?.size ?: 0
+                    if (backstackCount < lastKeyCount) {
+                        // The set of the ids associated with this key
+                        @Suppress("PrimitiveInCollection") // The order of the element matters
+                        val idsSet = localInfo.refCount[key]!!
+                        val id = idsSet.last()
+                        idsSet.remove(id)
+                        if (!localInfo.idsInComposition.contains(id)) {
+                            localInfo.savedStateHolder!!.removeState(id)
+                        }
                     }
                     // If the refCount is 0, remove the key from the refCount.
-                    if (localInfo.refCount[key] == 0) {
+                    if (localInfo.refCount[key]?.isEmpty() == true) {
                         localInfo.refCount.remove(key)
                     }
                 }
@@ -80,8 +84,46 @@
     public override fun <T : Any> ProvideToEntry(entry: NavEntry<T>) {
         val localInfo = LocalSaveableStateNavLocalInfo.current
         val key = entry.key
-        val refCount = localInfo.refCount[key]
-        val id: Int = rememberSaveable(key, refCount) { getIdForKey(key, refCount) }
+        // Tracks whether the key is changed
+        var keyChanged = false
+        var id: Int =
+            rememberSaveable(key) {
+                keyChanged = true
+                localInfo.refCount[key]!!.last()
+            }
+        id =
+            rememberSaveable(localInfo.refCount[key]?.size) {
+                // if the key changed, use the current id
+                // If the key was not changed, and the current id is not in composition or on the
+                // back
+                // stack then update the id with the last item from the backstack with the
+                // associated
+                // key. This ensures that we can handle duplicates, both consecutive and
+                // non-consecutive
+                if (
+                    !keyChanged &&
+                        (!localInfo.idsInComposition.contains(id) ||
+                            localInfo.refCount[key]?.contains(id) == true)
+                ) {
+                    localInfo.refCount[key]!!.last()
+                } else {
+                    id
+                }
+            }
+        keyChanged = false
+        DisposableEffect(key1 = key) {
+            localInfo.idsInComposition.add(id)
+            onDispose {
+                if (localInfo.idsInComposition.remove(id) && !localInfo.refCount.contains(key)) {
+                    localInfo.savedStateHolder!!.removeState(id)
+                    // If the refCount is 0, remove the key from the refCount.
+                    if (localInfo.refCount[key]?.isEmpty() == true) {
+                        localInfo.refCount.remove(key)
+                    }
+                }
+            }
+        }
+
         localInfo.savedStateHolder?.SaveableStateProvider(id) { entry.content.invoke(key) }
     }
 }
@@ -96,7 +138,9 @@
 
 internal class SaveableStateNavLocalInfo {
     internal var savedStateHolder: SaveableStateHolder? = null
-    internal val refCount: MutableObjectIntMap<Any> = MutableObjectIntMap()
+    internal val refCount: MutableMap<Any, LinkedHashSet<Int>> = mutableMapOf()
+    @Suppress("PrimitiveInCollection") // The order of the element matters
+    internal val idsInComposition: LinkedHashSet<Int> = LinkedHashSet<Int>()
 }
 
 internal fun getIdForKey(key: Any, count: Int): Int = 31 * key.hashCode() + count
diff --git a/pdf/pdf-viewer-fragment/src/main/java/androidx/pdf/viewer/fragment/PdfViewerFragment.kt b/pdf/pdf-viewer-fragment/src/main/java/androidx/pdf/viewer/fragment/PdfViewerFragment.kt
index d96109b..ac5532a 100644
--- a/pdf/pdf-viewer-fragment/src/main/java/androidx/pdf/viewer/fragment/PdfViewerFragment.kt
+++ b/pdf/pdf-viewer-fragment/src/main/java/androidx/pdf/viewer/fragment/PdfViewerFragment.kt
@@ -684,7 +684,12 @@
         paginatedView?.selectionModel = updatedSelectionModel
 
         selectionActionMode =
-            SelectionActionMode(requireActivity(), paginatedView!!, updatedSelectionModel)
+            SelectionActionMode(
+                requireActivity(),
+                paginatedView!!,
+                zoomView!!,
+                updatedSelectionModel
+            )
         selectionHandles =
             PdfSelectionHandles(
                 updatedSelectionModel,
diff --git a/pdf/pdf-viewer/src/main/java/androidx/pdf/select/SelectionActionMode.java b/pdf/pdf-viewer/src/main/java/androidx/pdf/select/SelectionActionMode.java
index dc4ebd0..1aeda00e 100644
--- a/pdf/pdf-viewer/src/main/java/androidx/pdf/select/SelectionActionMode.java
+++ b/pdf/pdf-viewer/src/main/java/androidx/pdf/select/SelectionActionMode.java
@@ -31,12 +31,9 @@
 import androidx.pdf.R;
 import androidx.pdf.models.PageSelection;
 import androidx.pdf.models.SelectionBoundary;
-import androidx.pdf.util.ObservableValue;
 import androidx.pdf.util.Preconditions;
-import androidx.pdf.viewer.PageMosaicView;
-import androidx.pdf.viewer.PageViewFactory;
 import androidx.pdf.viewer.PaginatedView;
-import androidx.pdf.viewer.PdfSelectionHandles;
+import androidx.pdf.widget.ZoomView;
 
 import org.jspecify.annotations.NonNull;
 
@@ -47,6 +44,7 @@
 public class SelectionActionMode {
     private static final String TAG = "SelectionActionMode";
     private final Context mContext;
+    private final ZoomView mZoomView;
     private final PaginatedView mPaginatedView;
     private final SelectionModel<PageSelection> mSelectionModel;
     private final Object mSelectionObserverKey;
@@ -57,30 +55,28 @@
     private final String mKeyCopiedText = "PdfCopiedText";
 
     public SelectionActionMode(@NonNull Context context, @NonNull PaginatedView paginatedView,
-            @NonNull SelectionModel<PageSelection> selectionModel) {
+            @NonNull ZoomView zoomView, @NonNull SelectionModel<PageSelection> selectionModel) {
         Preconditions.checkNotNull(context, "Context should not be null");
         Preconditions.checkNotNull(paginatedView, "paginatedView should not be null");
         Preconditions.checkNotNull(paginatedView, "selectionModel should not be null");
         Preconditions.checkNotNull(paginatedView, "callback should not be null");
         this.mContext = context;
         this.mPaginatedView = paginatedView;
+        this.mZoomView = zoomView;
         this.mSelectionModel = selectionModel;
         this.mCallback = new SelectionCallback();
 
         mSelectionObserverKey = selectionModel.selection().addObserver(
-                new ObservableValue.ValueObserver<PageSelection>() {
-                    @Override
-                    public void onChange(PageSelection oldValue, PageSelection newValue) {
-                        mCurrentSelection = newValue;
-                        if (newValue == null) {
-                            stopActionMode();
-                        } else if (oldValue == null) {
-                            startActionMode();
-                        } else {
-                            if (!oldValue.getStart().equals(newValue.getStart())
-                                    && !oldValue.getStop().equals(newValue.getStop())) {
-                                resume();
-                            }
+                (oldValue, newValue) -> {
+                    mCurrentSelection = newValue;
+                    if (newValue == null) {
+                        stopActionMode();
+                    } else if (oldValue == null) {
+                        startActionMode();
+                    } else {
+                        if (!oldValue.getStart().equals(newValue.getStart())
+                                && !oldValue.getStop().equals(newValue.getStop())) {
+                            resume();
                         }
                     }
                 });
@@ -95,12 +91,9 @@
 
     /** Start this action mode - updates the menu, ensures it is visible. */
     private void startActionMode() {
-        PageViewFactory.PageView pageView = mPaginatedView.getViewAt(
-                Objects.requireNonNull(mSelectionModel.mSelection.get()).getPage());
-        if (pageView != null) {
-            PageMosaicView pageMosaicView = pageView.getPageView();
-            pageMosaicView.startActionMode(mCallback, ActionMode.TYPE_FLOATING);
-        }
+        mZoomView.post(() -> {
+            mZoomView.startActionMode(mCallback, ActionMode.TYPE_FLOATING);
+        });
     }
 
     /** Resumes the context menu **/
@@ -168,7 +161,6 @@
         @Override
         public void onDestroyActionMode(ActionMode mode) {
             mode = null;
-
         }
 
         /**
@@ -177,40 +169,66 @@
          */
         @Override
         public void onGetContentRect(ActionMode mode, View view, Rect outRect) {
-            Rect bounds = getBoundsToPlaceMenu();
-            outRect.set(bounds.left, bounds.top, bounds.right, bounds.bottom);
+            Rect bounds = chooseContentLocationForMenu();
+            outRect.set(Math.round(mZoomView.toZoomViewX(bounds.left)),
+                    Math.round(mZoomView.toZoomViewY(bounds.top)),
+                    Math.round(mZoomView.toZoomViewX(bounds.right)),
+                    Math.round(mZoomView.toZoomViewY(bounds.bottom)));
         }
 
-        private Rect getBoundsToPlaceMenu() {
-            PageSelection pageSelection = mSelectionModel.mSelection.get();
-            int selectionPage = pageSelection.getPage();
-            PdfSelectionHandles mSelectionHandles = mPaginatedView.getSelectionHandles();
-            Rect startHandlerect = new Rect();
-            mSelectionHandles.getStartHandle().getGlobalVisibleRect(startHandlerect);
+        /**
+         * Determines the ideal location of the selection menu based on the visible area of the PDF
+         * and the bounds of the current selection.
+         *
+         * In order of preference, this will return:
+         * a) The first selection boundary, if it's visible
+         * b) The last selection boundary, if the first is not visible
+         * c) The middle of the visible area, if neither the first nor the last selection boundary
+         * are visible
+         *
+         * @return the ideal location for the selection menu in content coordinates, as a
+         * {@link Rect}. These coordinates should be adjusted for the current zoom / scroll level
+         * before using them in any {@link View}
+         */
+        private Rect chooseContentLocationForMenu() {
+            PageSelection selection = Objects.requireNonNull(mSelectionModel.mSelection.get());
+            int selectionPage = selection.getPage();
+            Rect pageLocation = mPaginatedView.getModel().getPageLocation(selectionPage,
+                    mPaginatedView.getViewArea());
 
-            Rect stopHandleRect = new Rect();
-            mSelectionHandles.getStopHandle().getGlobalVisibleRect(stopHandleRect);
+            List<Rect> selectionBounds = selection.getRects();
+            // The selection bounds are defined in page coordinates, we need their bounds in the
+            // overall PDF.
+            Rect firstSelectionBounds = toContentRect(selectionBounds.get(0), pageLocation);
+            Rect lastSelectionBounds = toContentRect(
+                    selectionBounds.get(selectionBounds.size() - 1), pageLocation);
 
-            int screenWidth = mPaginatedView.getResources().getDisplayMetrics().widthPixels;
-            int screenHeight = mPaginatedView.getResources().getDisplayMetrics().heightPixels;
-
-            if (pageSelection.getRects().size() == 1 || startHandlerect.intersect(0, 0, screenWidth,
-                    screenHeight)) {
-                return pageSelection.getRects().get(0);
-            } else if (stopHandleRect.intersect(0, 0, screenWidth, screenHeight)) {
-                List<Rect> rects = pageSelection.getRects();
-                return rects.get(rects.size() - 1);
+            Rect visibleArea = mZoomView.getVisibleAreaInContentCoords();
+            if (firstSelectionBounds.intersect(visibleArea)) {
+                return firstSelectionBounds;
+            } else if (lastSelectionBounds.intersect(visibleArea)) {
+                return lastSelectionBounds;
             } else {
-                // Center of the view in page coordinates
-                int viewCentreX = mPaginatedView.getViewArea().centerX()
-                        * mPaginatedView.getModel().getPageSize(selectionPage).getWidth()
-                        / mPaginatedView.getModel().getWidth();
-                int viewCentreY = mPaginatedView.getViewArea().centerY()
-                        - mPaginatedView.getModel().getPageLocation(selectionPage,
-                        mPaginatedView.getViewArea()).top;
-                return new Rect(viewCentreX, viewCentreY, viewCentreX, viewCentreY);
+                // If neither the beginning nor the end of the current selection are visible,
+                // center the menu in View
+                int centerX = visibleArea.centerX();
+                int centerY = visibleArea.centerY();
+                return new Rect(centerX - 1, centerY - 1, centerX + 1, centerY + 1);
             }
         }
+
+        /**
+         * Converts a {@link Rect} in PDF page coordinate space to a Rect in PDF coordinate space
+         *
+         * @param pageRect     A {@link Rect} describing a location within a PDF page
+         * @param pageLocation A {@link Rect} describing the location of a page within a PDF
+         * @return A {@link Rect} describing {@code pageRect}'s location in the overall PDF
+         */
+        private Rect toContentRect(Rect pageRect, Rect pageLocation) {
+            Rect out = new Rect(pageRect);
+            out.offset(pageLocation.left, pageLocation.top);
+            return out;
+        }
     }
 }
 
diff --git a/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/ZoomScrollValueObserver.java b/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/ZoomScrollValueObserver.java
index 94a355e..b1e5098 100644
--- a/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/ZoomScrollValueObserver.java
+++ b/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/ZoomScrollValueObserver.java
@@ -33,6 +33,8 @@
 import org.jspecify.annotations.NonNull;
 import org.jspecify.annotations.Nullable;
 
+import java.util.List;
+
 @RestrictTo(RestrictTo.Scope.LIBRARY)
 public class ZoomScrollValueObserver implements ObservableValue.ValueObserver<ZoomView.ZoomScroll> {
     private final PaginatedView mPaginatedView;
@@ -140,19 +142,36 @@
             // If selection is within the range of visible pages
             if (selectionPage >= firstPageInVisibleRange
                     && selectionPage <= lastPageInVisisbleRange) {
+                List<Rect> selectionRects =
+                        mPaginatedView.getSelectionModel().selection().get().getRects();
+                int startX = Integer.MAX_VALUE;
+                int startY = Integer.MAX_VALUE;
+                int endX = Integer.MIN_VALUE;
+                int endY = Integer.MIN_VALUE;
+                for (Rect rect : selectionRects) {
+                    if (rect.left < startX) {
+                        startX = rect.left;
+                    }
+                    if (rect.top < startY) {
+                        startY = rect.top;
+                    }
+                    if (rect.right > endX) {
+                        endX = rect.right;
+                    }
+                    if (rect.bottom > endY) {
+                        endY = rect.bottom;
+                    }
+                }
+
                 // Start and stop coordinates in a page wrt pagination model
-                int startX = mPaginatedView.getModel().getLookAtX(selectionPage,
-                        mPaginatedView.getSelectionModel().selection().get().getStart().getX());
-                int startY = mPaginatedView.getModel().getLookAtY(selectionPage,
-                        mPaginatedView.getSelectionModel().selection().get().getStart().getY());
-                int stopX = mPaginatedView.getModel().getLookAtX(selectionPage,
-                        mPaginatedView.getSelectionModel().selection().get().getStop().getX());
-                int stopY = mPaginatedView.getModel().getLookAtY(selectionPage,
-                        mPaginatedView.getSelectionModel().selection().get().getStop().getY());
+                startX = mPaginatedView.getModel().getLookAtX(selectionPage, startX);
+                startY = mPaginatedView.getModel().getLookAtY(selectionPage, startY);
+                endX = mPaginatedView.getModel().getLookAtX(selectionPage, endX);
+                endY = mPaginatedView.getModel().getLookAtY(selectionPage, endY);
 
                 Rect currentViewArea = mPaginatedView.getViewArea();
 
-                if (currentViewArea.intersect(startX, startY, stopX, stopY)) {
+                if (currentViewArea.intersects(startX, startY, endX, endY)) {
                     mSelectionActionMode.resume();
                 }
             }
diff --git a/pdf/pdf-viewer/src/main/java/androidx/pdf/widget/ZoomView.java b/pdf/pdf-viewer/src/main/java/androidx/pdf/widget/ZoomView.java
index afe92c5..f88fbc5 100644
--- a/pdf/pdf-viewer/src/main/java/androidx/pdf/widget/ZoomView.java
+++ b/pdf/pdf-viewer/src/main/java/androidx/pdf/widget/ZoomView.java
@@ -1078,14 +1078,20 @@
     }
 
     /**
-     * Given a point in the content's co-ordinates, convert it to the zoom-view's co-ordinates,
-     * using the current zoom and scroll position of the zoomview.
+     * Given an X coordinate in content space, return the same X coordinate in this
+     * {@link ZoomView}'s coordinate space, accounting for the current {@link #getScrollX()} and
+     * {@link #getZoom()} properties.
      */
-    protected float toZoomViewX(float contentX) {
+    public float toZoomViewX(float contentX) {
         return ZoomUtils.toZoomViewCoordinate(contentX, getZoom(), getScrollX());
     }
 
-    protected float toZoomViewY(float contentY) {
+    /**
+     * Given a Y coordinate in content space, return the same Y coordinate in this
+     * {@link ZoomView}'s coordinate space, accounting for the current {@link #getScrollY()} and
+     * {@link #getZoom()} properties.
+     */
+    public float toZoomViewY(float contentY) {
         return ZoomUtils.toZoomViewCoordinate(contentY, getZoom(), getScrollY());
     }
 
diff --git a/pdf/pdf-viewer/src/main/kotlin/androidx/pdf/view/PdfView.kt b/pdf/pdf-viewer/src/main/kotlin/androidx/pdf/view/PdfView.kt
index 8fc6c31..9313a3f 100644
--- a/pdf/pdf-viewer/src/main/kotlin/androidx/pdf/view/PdfView.kt
+++ b/pdf/pdf-viewer/src/main/kotlin/androidx/pdf/view/PdfView.kt
@@ -1179,7 +1179,7 @@
 
             // Else, center the context menu in view
             val centerX = (pdfView.x + pdfView.width / 2).roundToInt()
-            val centerY = (pdfView.x + pdfView.height / 2).roundToInt()
+            val centerY = (pdfView.y + pdfView.height / 2).roundToInt()
             outRect.set(centerX, centerY, centerX + 1, centerY + 1)
         }
     }
diff --git a/pdf/pdf-viewer/src/test/java/androidx/pdf/select/SelectionActionModeTest.java b/pdf/pdf-viewer/src/test/java/androidx/pdf/select/SelectionActionModeTest.java
index 381a74e..3bb84af 100644
--- a/pdf/pdf-viewer/src/test/java/androidx/pdf/select/SelectionActionModeTest.java
+++ b/pdf/pdf-viewer/src/test/java/androidx/pdf/select/SelectionActionModeTest.java
@@ -17,26 +17,29 @@
 package androidx.pdf.select;
 
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.content.Context;
+import android.view.ActionMode;
 
 import androidx.pdf.models.PageSelection;
 import androidx.pdf.util.ObservableValue;
 import androidx.pdf.util.Observables;
-import androidx.pdf.viewer.PageMosaicView;
-import androidx.pdf.viewer.PageViewFactory;
 import androidx.pdf.viewer.PaginatedView;
+import androidx.pdf.widget.ZoomView;
 
 import org.jspecify.annotations.NonNull;
 import org.jspecify.annotations.Nullable;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
 
+@RunWith(RobolectricTestRunner.class)
 public class SelectionActionModeTest {
 
     @Mock
@@ -51,16 +54,13 @@
     private PaginatedView mPaginatedView;
 
     @Mock
-    private PageMosaicView mPageMosaicView;
-
-    @Mock
     private PageSelection mPageSelection;
 
     @Mock
-    Observables.ExposedValue<PageSelection> mSelection;
+    private Observables.ExposedValue<PageSelection> mSelection;
 
     @Mock
-    PageViewFactory.PageView mPageView;
+    private ZoomView mZoomView;
 
 
     @Before
@@ -70,6 +70,11 @@
 
     @Test
     public void testStartActionMode() {
+        when(mZoomView.post(any())).thenAnswer(invocation -> {
+            Runnable action = invocation.getArgument(0);
+            action.run();
+            return true;
+        });
         SelectionModel<PageSelection> selectionModel = new SelectionModel<PageSelection>() {
 
             @Override
@@ -102,20 +107,18 @@
 
         selectionModel.setSelection(mPageSelection);
 
-        when(mPaginatedView.getViewAt(anyInt())).thenReturn(mPageView);
-        when(mPageView.getPageView()).thenReturn(mPageMosaicView);
+        mSelectionActionMode = new SelectionActionMode(mContext, mPaginatedView, mZoomView,
+                selectionModel);
 
-        mSelectionActionMode = new SelectionActionMode(mContext, mPaginatedView, selectionModel);
-
-        verify(mPaginatedView).getViewAt(anyInt());
-        verify(mPageMosaicView).startActionMode(any(), anyInt());
+        verify(mZoomView).startActionMode(any(), eq(ActionMode.TYPE_FLOATING));
     }
 
     @Test
     public void testDestroyRemoveObserver() {
         when(mSelectionModel.selection()).thenReturn(mSelection);
 
-        mSelectionActionMode = new SelectionActionMode(mContext, mPaginatedView, mSelectionModel);
+        mSelectionActionMode = new SelectionActionMode(mContext, mPaginatedView, mZoomView,
+                mSelectionModel);
 
         mSelectionActionMode.destroy();
 
diff --git a/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/ZoomScrollValueObserverTest.java b/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/ZoomScrollValueObserverTest.java
index a3420d2..807b30f 100644
--- a/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/ZoomScrollValueObserverTest.java
+++ b/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/ZoomScrollValueObserverTest.java
@@ -29,7 +29,6 @@
 import androidx.pdf.data.Range;
 import androidx.pdf.find.FindInFileView;
 import androidx.pdf.models.PageSelection;
-import androidx.pdf.models.SelectionBoundary;
 import androidx.pdf.select.SelectionActionMode;
 import androidx.pdf.util.ObservableValue;
 import androidx.pdf.util.Observables;
@@ -45,6 +44,9 @@
 import org.junit.runner.RunWith;
 import org.robolectric.RobolectricTestRunner;
 
+import java.util.ArrayList;
+import java.util.List;
+
 @SmallTest
 @RunWith(RobolectricTestRunner.class)
 public class ZoomScrollValueObserverTest {
@@ -102,8 +104,10 @@
 
             }
         });
-        when(mMockPageSelection.getStart()).thenReturn(new SelectionBoundary(0, 0, 0, false));
-        when(mMockPageSelection.getStop()).thenReturn(new SelectionBoundary(0, 100, 100, false));
+        List<Rect> selectionBoundaries = new ArrayList<>();
+        selectionBoundaries.add(new Rect(0, 0, 10, 10));
+        selectionBoundaries.add(new Rect(90, 90, 100, 100));
+        when(mMockPageSelection.getRects()).thenReturn(selectionBoundaries);
         when(mMockPaginatedView.getViewArea()).thenReturn(RECT);
         when(mMockPaginationModel.getLookAtX(0, 0)).thenReturn(1);
         when(mMockPaginationModel.getLookAtX(0, 100)).thenReturn(50);
diff --git a/privacysandbox/activity/OWNERS b/privacysandbox/activity/OWNERS
new file mode 100644
index 0000000..87d7833
--- /dev/null
+++ b/privacysandbox/activity/OWNERS
@@ -0,0 +1,5 @@
[email protected]
[email protected]
[email protected]
[email protected]
[email protected]
\ No newline at end of file
diff --git a/privacysandbox/ui/integration-tests/testapp/build.gradle b/privacysandbox/ui/integration-tests/testapp/build.gradle
index e88938e..b3b181e 100644
--- a/privacysandbox/ui/integration-tests/testapp/build.gradle
+++ b/privacysandbox/ui/integration-tests/testapp/build.gradle
@@ -72,4 +72,7 @@
     implementation(project(":privacysandbox:ui:integration-tests:testsdkproviderwrapper"))
     implementation(project(":privacysandbox:ui:ui-client"))
     implementation(project(":privacysandbox:ui:ui-provider"))
+    implementation(libs.media3Ui)
+    implementation(libs.media3Common)
+    implementation(libs.media3Exoplayer)
 }
diff --git a/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/BaseFragment.kt b/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/BaseFragment.kt
index 690d4f7..e269009 100644
--- a/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/BaseFragment.kt
+++ b/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/BaseFragment.kt
@@ -34,6 +34,7 @@
 import androidx.privacysandbox.ui.integration.sdkproviderutils.SdkApiConstants.Companion.AdFormat
 import androidx.privacysandbox.ui.integration.sdkproviderutils.SdkApiConstants.Companion.AdType
 import androidx.privacysandbox.ui.integration.sdkproviderutils.SdkApiConstants.Companion.MediationOption
+import androidx.privacysandbox.ui.integration.testapp.util.AdHolder
 import androidx.privacysandbox.ui.integration.testsdkprovider.ISdkApi
 import androidx.privacysandbox.ui.integration.testsdkprovider.ISdkApiFactory
 import kotlinx.coroutines.CoroutineScope
@@ -118,6 +119,28 @@
         drawViewabilityLayer: Boolean
     ) {}
 
+    fun loadAd(
+        adHolder: AdHolder,
+        @AdFormat adFormat: Int,
+        @AdType adType: Int,
+        @MediationOption mediationOption: Int,
+        drawViewabilityLayer: Boolean,
+        waitInsideOnDraw: Boolean = false
+    ) {
+        CoroutineScope(Dispatchers.Main).launch {
+            val sdkBundle =
+                sdkApi.loadAd(
+                    adFormat,
+                    adType,
+                    mediationOption,
+                    waitInsideOnDraw,
+                    drawViewabilityLayer
+                )
+            adHolder.populateAd(sdkBundle, adFormat)
+        }
+    }
+
+    // TODO(b/369355774): replace with loadAd on all supported fragments
     fun loadBannerAd(
         @AdType adType: Int,
         @MediationOption mediationOption: Int,
diff --git a/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/MainActivity.kt b/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/MainActivity.kt
index 58eea59..b9c03b4 100644
--- a/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/MainActivity.kt
+++ b/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/MainActivity.kt
@@ -216,7 +216,12 @@
                     applicationContext,
                     resources.getStringArray(R.array.ad_format_menu_array)
                 ) { position: Int ->
-                    isSupportedOptionsCombination(position)
+                    val isSupported = isSupportedOptionsCombination(position, mediationOption)
+                    when (position) {
+                        AdFormat.BANNER_AD -> isSupported
+                        AdFormat.NATIVE_AD -> isSupported && supportsNativeAd(currentFragment)
+                        else -> false
+                    }
                 }
             onItemSelectedListener = OnItemSelectedListener()
         }
@@ -229,8 +234,8 @@
                 DisabledItemsArrayAdapter(
                     applicationContext,
                     resources.getStringArray(R.array.mediation_dropdown_menu_array)
-                ) { _: Int ->
-                    isSupportedOptionsCombination(adFormat)
+                ) { position: Int ->
+                    isSupportedOptionsCombination(adFormat, position)
                 }
             onItemSelectedListener = OnItemSelectedListener()
         }
@@ -243,7 +248,7 @@
                     applicationContext,
                     resources.getStringArray(R.array.ad_type_dropdown_menu_array)
                 ) { _: Int ->
-                    isSupportedOptionsCombination(adFormat)
+                    isSupportedOptionsCombination(adFormat, mediationOption)
                 }
             onItemSelectedListener = OnItemSelectedListener()
         }
@@ -348,20 +353,34 @@
 
     private fun isSupportedOptionsCombination(
         @AdFormat adFormat: Int,
+        @MediationOption mediationOption: Int
     ): Boolean {
         when (adFormat) {
             AdFormat.BANNER_AD -> return true
-            AdFormat.NATIVE_AD -> return false
+            AdFormat.NATIVE_AD -> {
+                when (mediationOption) {
+                    MediationOption.NON_MEDIATED,
+                    MediationOption.IN_APP_MEDIATEE,
+                    MediationOption.SDK_RUNTIME_MEDIATEE -> return true
+                    MediationOption.SDK_RUNTIME_MEDIATEE_WITH_OVERLAY,
+                    MediationOption.REFRESHABLE_MEDIATION -> return false
+                }
+            }
         }
         return false
     }
 
+    // TODO(b/3300859): remove once all non-fullscreen fragments are supported for native.
+    private fun supportsNativeAd(fragment: BaseFragment): Boolean = fragment is ResizeFragment
+
     private fun updateDrawerOptions() {
         setAllControlsEnabled(true)
         if (adFormat == AdFormat.NATIVE_AD) {
-            runOnUiThread { navigationView.menu.forEach { it.isEnabled = false } }
+            runOnUiThread {
+                navigationView.menu.findItem(R.id.item_scroll).isEnabled = false
+                navigationView.menu.findItem(R.id.item_pooling_container).isEnabled = false
+            }
             viewabilityToggleButton.isEnabled = false
-            zOrderToggleButton.isEnabled = false
             composeToggleButton.isEnabled = false
         }
     }
@@ -377,7 +396,7 @@
     }
 
     private fun switchContentFragment(fragment: BaseFragment, title: CharSequence?): Boolean {
-        setAllControlsEnabled(true)
+        updateDrawerOptions()
         drawerLayout.closeDrawers()
         supportFragmentManager
             .beginTransaction()
diff --git a/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/ResizeFragment.kt b/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/ResizeFragment.kt
index 260f1fd..3d5ebcf4 100644
--- a/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/ResizeFragment.kt
+++ b/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/ResizeFragment.kt
@@ -16,27 +16,33 @@
 
 package androidx.privacysandbox.ui.integration.testapp
 
+import android.graphics.Color
 import android.os.Bundle
+import android.util.TypedValue
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import android.view.ViewGroup.MarginLayoutParams
 import android.widget.Button
+import androidx.core.view.setMargins
 import androidx.privacysandbox.ui.client.view.SandboxedSdkView
 import androidx.privacysandbox.ui.integration.sdkproviderutils.SdkApiConstants.Companion.AdFormat
 import androidx.privacysandbox.ui.integration.sdkproviderutils.SdkApiConstants.Companion.AdType
 import androidx.privacysandbox.ui.integration.sdkproviderutils.SdkApiConstants.Companion.MediationOption
+import androidx.privacysandbox.ui.integration.testapp.util.AdHolder
 import kotlin.math.max
 
 class ResizeFragment : BaseFragment() {
 
-    private lateinit var resizableBannerView: SandboxedSdkView
+    private lateinit var resizableAdHolder: AdHolder
+
     private lateinit var resizeButton: Button
     private lateinit var resizeFromSdkButton: Button
     private lateinit var setPaddingButton: Button
     private lateinit var inflatedView: View
 
     override fun getSandboxedSdkViews(): List<SandboxedSdkView> {
-        return listOf(resizableBannerView)
+        return resizableAdHolder.sandboxedSdkViews
     }
 
     override fun handleLoadAdFromDrawer(
@@ -49,15 +55,14 @@
         currentAdType = adType
         currentMediationOption = mediationOption
         shouldDrawViewabilityLayer = drawViewabilityLayer
-        if (adFormat == AdFormat.BANNER_AD) {
-            loadBannerAd(
-                adType,
-                mediationOption,
-                resizableBannerView,
-                drawViewabilityLayer,
-                waitInsideOnDraw = true
-            )
-        }
+        loadAd(
+            resizableAdHolder,
+            currentAdFormat,
+            currentAdType,
+            currentMediationOption,
+            shouldDrawViewabilityLayer,
+            waitInsideOnDraw = true
+        )
     }
 
     override fun onCreateView(
@@ -66,25 +71,32 @@
         savedInstanceState: Bundle?
     ): View {
         inflatedView = inflater.inflate(R.layout.fragment_resize, container, false)
-        resizableBannerView = inflatedView.findViewById(R.id.resizable_ad_view)
+        resizableAdHolder =
+            inflatedView.findViewById<AdHolder>(R.id.resizable_ad_view).apply {
+                adViewLayoutParams =
+                    MarginLayoutParams(adViewLayoutParams).apply {
+                        setMargins(convertFromDpToPixels(MARGIN_DP))
+                    }
+                adViewBackgroundColor = Color.parseColor(AD_VIEW_BACKGROUND_COLOR)
+            }
         resizeButton = inflatedView.findViewById(R.id.resize_button)
         resizeFromSdkButton = inflatedView.findViewById(R.id.resize_sdk_button)
         setPaddingButton = inflatedView.findViewById(R.id.set_padding_button)
-        if (currentAdFormat == AdFormat.BANNER_AD) {
-            loadResizableBannerAd()
-        }
+        initResizeButton()
+        initSetPaddingButton()
+
+        loadAd(
+            resizableAdHolder,
+            currentAdFormat,
+            currentAdType,
+            currentMediationOption,
+            shouldDrawViewabilityLayer,
+            true
+        )
         return inflatedView
     }
 
-    private fun loadResizableBannerAd() {
-        loadBannerAd(
-            currentAdType,
-            currentMediationOption,
-            resizableBannerView,
-            shouldDrawViewabilityLayer,
-            waitInsideOnDraw = true
-        )
-
+    private fun initResizeButton() {
         val displayMetrics = resources.displayMetrics
         val maxSizePixels = Math.min(displayMetrics.widthPixels, displayMetrics.heightPixels)
 
@@ -93,21 +105,24 @@
         }
 
         resizeButton.setOnClickListener {
-            val newWidth = newSize(resizableBannerView.width, maxSizePixels)
-            val newHeight = newSize(resizableBannerView.height, maxSizePixels)
-            resizableBannerView.layoutParams =
-                resizableBannerView.layoutParams.apply {
+            val newWidth = newSize(resizableAdHolder.currentAdView.width, maxSizePixels)
+            val newHeight =
+                newSize(resizableAdHolder.currentAdView.height, resizableAdHolder.height)
+            resizableAdHolder.currentAdView.layoutParams =
+                resizableAdHolder.currentAdView.layoutParams.apply {
                     width = newWidth
                     height = newHeight
                 }
         }
+    }
 
+    private fun initSetPaddingButton() {
         setPaddingButton.setOnClickListener {
             // Set halfWidth and halfHeight to minimum 10 to avoid crashes when the width and height
             // are very small
-            val halfWidth = max(10, (resizableBannerView.width / 2) - 10)
-            val halfHeight = max(10, resizableBannerView.height / 2) - 10
-            resizableBannerView.setPadding(
+            val halfWidth = max(10, (resizableAdHolder.currentAdView.width / 2) - 10)
+            val halfHeight = max(10, (resizableAdHolder.currentAdView.height / 2) - 10)
+            resizableAdHolder.currentAdView.setPadding(
                 (10..halfWidth).random(),
                 (10..halfHeight).random(),
                 (10..halfWidth).random(),
@@ -115,4 +130,17 @@
             )
         }
     }
+
+    private fun convertFromDpToPixels(dpValue: Float): Int =
+        TypedValue.applyDimension(
+                TypedValue.COMPLEX_UNIT_DIP,
+                dpValue,
+                context?.resources?.displayMetrics
+            )
+            .toInt()
+
+    private companion object {
+        const val MARGIN_DP = 16.0f
+        const val AD_VIEW_BACKGROUND_COLOR = "#D3D3D3"
+    }
 }
diff --git a/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/util/AdHolder.kt b/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/util/AdHolder.kt
new file mode 100644
index 0000000..faebc27
--- /dev/null
+++ b/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/util/AdHolder.kt
@@ -0,0 +1,94 @@
+/*
+ * Copyright 2025 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.ui.integration.testapp.util
+
+import android.annotation.SuppressLint
+import android.content.Context
+import android.graphics.Color
+import android.os.Bundle
+import android.util.AttributeSet
+import android.view.View
+import android.view.ViewGroup
+import android.widget.FrameLayout
+import android.widget.TextView
+import androidx.privacysandbox.ui.client.SandboxedUiAdapterFactory
+import androidx.privacysandbox.ui.client.view.SandboxedSdkView
+import androidx.privacysandbox.ui.client.view.SharedUiContainer
+import androidx.privacysandbox.ui.core.ExperimentalFeatures
+import androidx.privacysandbox.ui.integration.sdkproviderutils.SdkApiConstants.Companion.AdFormat
+import androidx.privacysandbox.ui.integration.testapp.R
+
+@SuppressLint("NullAnnotationGroup")
+@OptIn(ExperimentalFeatures.SharedUiPresentationApi::class)
+class AdHolder(context: Context, attrs: AttributeSet? = null) : FrameLayout(context, attrs) {
+    private val nativeAdLoader = NativeAdLoader(context)
+    private val bannerAdView: SandboxedSdkView = SandboxedSdkView(context)
+    private val nativeAdView: SharedUiContainer
+        get() = nativeAdLoader.adView
+
+    val sandboxedSdkViews: List<SandboxedSdkView>
+        get() =
+            when (currentAdFormat) {
+                AdFormat.BANNER_AD -> listOf(bannerAdView)
+                AdFormat.NATIVE_AD ->
+                    listOf(
+                        nativeAdView.findViewById(R.id.native_ad_remote_overlay_icon),
+                        nativeAdView.findViewById(R.id.native_ad_media_view_1)
+                    )
+                else -> listOf()
+            }
+
+    private var _currentAdFormat = AdFormat.BANNER_AD
+    var currentAdFormat: Int
+        get() = _currentAdFormat
+        private set(value) {
+            _currentAdFormat = value
+            currentAdView =
+                when (currentAdFormat) {
+                    AdFormat.BANNER_AD -> bannerAdView
+                    AdFormat.NATIVE_AD -> nativeAdView
+                    else ->
+                        TextView(context).apply {
+                            text = "Unsupported ad format."
+                            setTextColor(Color.RED)
+                        }
+                }
+        }
+
+    var currentAdView: View = bannerAdView
+        private set
+
+    var adViewLayoutParams: ViewGroup.LayoutParams =
+        ViewGroup.LayoutParams(LayoutParams.MATCH_PARENT, LayoutParams.MATCH_PARENT)
+    var adViewBackgroundColor: Int = Color.WHITE
+
+    fun populateAd(sdkBundle: Bundle, @AdFormat adFormat: Int) {
+        currentAdFormat = adFormat
+        removeAllViews()
+        when (adFormat) {
+            AdFormat.BANNER_AD -> populateBannerAd(sdkBundle)
+            AdFormat.NATIVE_AD -> nativeAdLoader.populateAd(sdkBundle)
+        }
+        currentAdView.layoutParams = adViewLayoutParams
+        currentAdView.setBackgroundColor(adViewBackgroundColor)
+        addView(currentAdView)
+    }
+
+    private fun populateBannerAd(sdkBundle: Bundle) {
+        bannerAdView.setAdapter(SandboxedUiAdapterFactory.createFromCoreLibInfo(sdkBundle))
+    }
+}
diff --git a/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/util/NativeAdLoader.kt b/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/util/NativeAdLoader.kt
new file mode 100644
index 0000000..55571a9
--- /dev/null
+++ b/privacysandbox/ui/integration-tests/testapp/src/main/java/androidx/privacysandbox/ui/integration/testapp/util/NativeAdLoader.kt
@@ -0,0 +1,154 @@
+/*
+ * Copyright 2025 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.ui.integration.testapp.util
+
+import android.annotation.SuppressLint
+import android.content.Context
+import android.graphics.BitmapFactory
+import android.graphics.Color
+import android.os.Bundle
+import android.view.View.inflate
+import android.widget.Button
+import android.widget.ImageView
+import android.widget.TextView
+import androidx.media3.ui.PlayerView
+import androidx.privacysandbox.ui.client.SandboxedUiAdapterFactory
+import androidx.privacysandbox.ui.client.SharedUiAdapterFactory
+import androidx.privacysandbox.ui.client.view.SandboxedSdkView
+import androidx.privacysandbox.ui.client.view.SharedUiAsset
+import androidx.privacysandbox.ui.client.view.SharedUiContainer
+import androidx.privacysandbox.ui.core.ExperimentalFeatures
+import androidx.privacysandbox.ui.integration.sdkproviderutils.PlayerViewProvider
+import androidx.privacysandbox.ui.integration.sdkproviderutils.SdkApiConstants.Companion.NativeAdAssetName
+import androidx.privacysandbox.ui.integration.sdkproviderutils.SdkApiConstants.Companion.NativeAdAssetProperties
+import androidx.privacysandbox.ui.integration.testapp.R
+
+@SuppressLint("NullAnnotationGroup")
+@OptIn(ExperimentalFeatures.SharedUiPresentationApi::class)
+class NativeAdLoader(context: Context) {
+    val adView: SharedUiContainer = inflate(context, NATIVE_AD_LAYOUT_ID, null) as SharedUiContainer
+    private val adHeadline: TextView = adView.findViewById(R.id.native_ad_headline)
+    private val adBody: TextView = adView.findViewById(R.id.native_ad_body)
+    private val adRemoteOverlayIcon: SandboxedSdkView =
+        adView.findViewById(R.id.native_ad_remote_overlay_icon)
+    private val adMediaView1: SandboxedSdkView = adView.findViewById(R.id.native_ad_media_view_1)
+    private val adOverlayIcon: ImageView = adView.findViewById(R.id.native_ad_overlay_icon)
+    private val adMediaView2: PlayerView = adView.findViewById(R.id.native_ad_media_view_2)
+    private val adCallToAction: Button = adView.findViewById(R.id.native_ad_call_to_action)
+
+    fun populateAd(sdkBundle: Bundle) {
+        adView.setAdapter(SharedUiAdapterFactory.createFromCoreLibInfo(sdkBundle))
+        val assets = sdkBundle.getBundle(NativeAdAssetName.ASSET_BUNDLE_NAME)
+
+        val headlineAssets = assets?.getBundle(NativeAdAssetName.HEADLINE)
+        adView.registerSharedUiAsset(
+            SharedUiAsset(
+                adHeadline.apply {
+                    text = headlineAssets?.getString(NativeAdAssetProperties.TEXT)
+                    setTextColor(
+                        Color.parseColor(headlineAssets?.getString(NativeAdAssetProperties.COLOR))
+                    )
+                },
+                NativeAdAssetName.HEADLINE
+            )
+        )
+
+        val bodyAssets = assets?.getBundle(NativeAdAssetName.BODY)
+        adView.registerSharedUiAsset(
+            SharedUiAsset(
+                adBody.apply {
+                    text = bodyAssets?.getString(NativeAdAssetProperties.TEXT)
+                    setTextColor(
+                        Color.parseColor(bodyAssets?.getString(NativeAdAssetProperties.COLOR))
+                    )
+                },
+                NativeAdAssetName.BODY
+            )
+        )
+
+        val adChoicesAssets = assets?.getBundle(NativeAdAssetName.AD_CHOICES)
+        if (adChoicesAssets != null) {
+            adView.registerSharedUiAsset(
+                SharedUiAsset(
+                    adRemoteOverlayIcon,
+                    NativeAdAssetName.AD_CHOICES,
+                    sandboxedUiAdapter =
+                        SandboxedUiAdapterFactory.createFromCoreLibInfo(adChoicesAssets)
+                )
+            )
+        }
+
+        val mediaView1Assets = assets?.getBundle(NativeAdAssetName.MEDIA_VIEW_1)
+        if (mediaView1Assets != null) {
+            adView.registerSharedUiAsset(
+                SharedUiAsset(
+                    adMediaView1,
+                    NativeAdAssetName.MEDIA_VIEW_1,
+                    sandboxedUiAdapter =
+                        SandboxedUiAdapterFactory.createFromCoreLibInfo(mediaView1Assets)
+                )
+            )
+        }
+
+        val iconAssets = assets?.getBundle(NativeAdAssetName.ICON)
+        adView.registerSharedUiAsset(
+            SharedUiAsset(
+                adOverlayIcon.apply {
+                    val iconByteArray = iconAssets?.getByteArray(NativeAdAssetProperties.BITMAP)!!
+                    val bitmap = BitmapFactory.decodeByteArray(iconByteArray, 0, iconByteArray.size)
+                    setImageBitmap(bitmap)
+                },
+                NativeAdAssetName.ICON
+            )
+        )
+
+        val mediaView2Assets = assets?.getBundle(NativeAdAssetName.MEDIA_VIEW_2)
+        adView.registerSharedUiAsset(
+            SharedUiAsset(
+                adMediaView2.apply {
+                    player =
+                        PlayerViewProvider()
+                            .PlayerWithState(
+                                context,
+                                mediaView2Assets?.getString(NativeAdAssetProperties.URL)!!
+                            )
+                            .initializePlayer()
+                },
+                NativeAdAssetName.MEDIA_VIEW_2
+            )
+        )
+
+        val callToActionAssets = assets?.getBundle(NativeAdAssetName.CALL_TO_ACTION)
+        adView.registerSharedUiAsset(
+            SharedUiAsset(
+                adCallToAction.apply {
+                    text = callToActionAssets?.getString(NativeAdAssetProperties.TEXT)
+                    setBackgroundColor(
+                        Color.parseColor(
+                            callToActionAssets?.getString(NativeAdAssetProperties.COLOR)
+                        )
+                    )
+                },
+                NativeAdAssetName.CALL_TO_ACTION
+            )
+        )
+    }
+
+    companion object {
+        val NATIVE_AD_LAYOUT_ID = R.layout.native_ad_layout
+    }
+}
diff --git a/privacysandbox/ui/integration-tests/testapp/src/main/res/layout/fragment_resize.xml b/privacysandbox/ui/integration-tests/testapp/src/main/res/layout/fragment_resize.xml
index 5c86f79..f52088d 100644
--- a/privacysandbox/ui/integration-tests/testapp/src/main/res/layout/fragment_resize.xml
+++ b/privacysandbox/ui/integration-tests/testapp/src/main/res/layout/fragment_resize.xml
@@ -21,16 +21,12 @@
     android:layout_weight="4"
     android:orientation="vertical">
 
-    <androidx.privacysandbox.ui.client.view.SandboxedSdkView
+    <androidx.privacysandbox.ui.integration.testapp.util.AdHolder
         android:id="@+id/resizable_ad_view"
         android:layout_width="wrap_content"
         android:layout_weight="2"
         android:layout_height="0dp"
-        android:layout_marginBottom="16dp"
-        android:layout_marginEnd="16dp"
-        android:layout_marginStart="16dp"
-        android:layout_marginTop="16dp"
-        android:background="#D3D3D3" />
+        android:background="#FFFFFF" />
 
     <LinearLayout
         android:layout_width="wrap_content"
diff --git a/privacysandbox/ui/integration-tests/testapp/src/main/res/layout/native_ad_layout.xml b/privacysandbox/ui/integration-tests/testapp/src/main/res/layout/native_ad_layout.xml
new file mode 100644
index 0000000..63a477e
--- /dev/null
+++ b/privacysandbox/ui/integration-tests/testapp/src/main/res/layout/native_ad_layout.xml
@@ -0,0 +1,84 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  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.
+  -->
+
+<androidx.privacysandbox.ui.client.view.SharedUiContainer xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/native_ad_container"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent">
+    <LinearLayout
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:weightSum="6"
+        android:orientation="vertical">
+        <!-- Ad Title, Ad Body and SandboxedSdkView overlay -->
+        <FrameLayout
+            android:layout_width="match_parent"
+            android:layout_weight="1"
+            android:layout_height="0dp">
+            <LinearLayout
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:orientation="vertical">
+                <TextView
+                    android:id="@+id/native_ad_headline"
+                    android:layout_width="wrap_content"
+                    android:layout_height="wrap_content"
+                    android:text="Headline not loaded"/>
+                <TextView
+                    android:id="@+id/native_ad_body"
+                    android:layout_width="wrap_content"
+                    android:layout_height="wrap_content"
+                    android:text="Body not loaded"/>
+            </LinearLayout>
+            <androidx.privacysandbox.ui.client.view.SandboxedSdkView
+                android:id="@+id/native_ad_remote_overlay_icon"
+                android:layout_width="50dp"
+                android:layout_height="50dp"
+                android:layout_gravity="right"/>
+        </FrameLayout>
+        <!-- MediaView #1 (SandboxedSdkView) and app-owned overlay icon -->
+        <FrameLayout
+            android:layout_width="match_parent"
+            android:layout_weight="2"
+            android:layout_height="0dp">
+            <androidx.privacysandbox.ui.client.view.SandboxedSdkView
+                android:id="@+id/native_ad_media_view_1"
+                android:layout_width="match_parent"
+                android:layout_height="match_parent"/>
+            <ImageView
+                android:id="@+id/native_ad_overlay_icon"
+                android:layout_width="50dp"
+                android:layout_height="50dp"
+                android:adjustViewBounds="true"/>
+        </FrameLayout>
+        <androidx.media3.ui.PlayerView
+            android:id="@+id/native_ad_media_view_2"
+            android:layout_width="match_parent"
+            android:layout_weight="2"
+            android:layout_height="0dp"/>
+        <FrameLayout
+            android:layout_width="match_parent"
+            android:layout_weight="1"
+            android:layout_height="0dp">
+            <Button
+                android:id="@+id/native_ad_call_to_action"
+                android:layout_width="match_parent"
+                android:layout_height="match_parent"
+                android:text="Call to action button not loaded"/>
+        </FrameLayout>
+    </LinearLayout>
+</androidx.privacysandbox.ui.client.view.SharedUiContainer>
diff --git a/savedstate/savedstate-compose/api/current.txt b/savedstate/savedstate-compose/api/current.txt
index de12ce4..ecb7dca 100644
--- a/savedstate/savedstate-compose/api/current.txt
+++ b/savedstate/savedstate-compose/api/current.txt
@@ -22,5 +22,29 @@
     method public static inline <reified T> androidx.savedstate.compose.serialization.serializers.MutableStateSerializer<T> MutableStateSerializer();
   }
 
+  public final class SnapshotStateListSerializer<T> implements kotlinx.serialization.KSerializer<androidx.compose.runtime.snapshots.SnapshotStateList<T>> {
+    ctor public SnapshotStateListSerializer(kotlinx.serialization.KSerializer<T> elementSerializer);
+    method public androidx.compose.runtime.snapshots.SnapshotStateList<T> deserialize(kotlinx.serialization.encoding.Decoder decoder);
+    method public kotlinx.serialization.descriptors.SerialDescriptor getDescriptor();
+    method public void serialize(kotlinx.serialization.encoding.Encoder encoder, androidx.compose.runtime.snapshots.SnapshotStateList<T> value);
+    property public kotlinx.serialization.descriptors.SerialDescriptor descriptor;
+  }
+
+  public final class SnapshotStateListSerializerKt {
+    method public static inline <reified T> androidx.savedstate.compose.serialization.serializers.SnapshotStateListSerializer<T> SnapshotStateListSerializer();
+  }
+
+  public final class SnapshotStateMapSerializer<K, V> implements kotlinx.serialization.KSerializer<androidx.compose.runtime.snapshots.SnapshotStateMap<K,V>> {
+    ctor public SnapshotStateMapSerializer(kotlinx.serialization.KSerializer<K> keySerializer, kotlinx.serialization.KSerializer<V> valueSerializer);
+    method public androidx.compose.runtime.snapshots.SnapshotStateMap<K,V> deserialize(kotlinx.serialization.encoding.Decoder decoder);
+    method public kotlinx.serialization.descriptors.SerialDescriptor getDescriptor();
+    method public void serialize(kotlinx.serialization.encoding.Encoder encoder, androidx.compose.runtime.snapshots.SnapshotStateMap<K,V> value);
+    property public kotlinx.serialization.descriptors.SerialDescriptor descriptor;
+  }
+
+  public final class SnapshotStateMapSerializerKt {
+    method public static inline <reified K, reified V> androidx.savedstate.compose.serialization.serializers.SnapshotStateMapSerializer<K,V> SnapshotStateMapSerializer();
+  }
+
 }
 
diff --git a/savedstate/savedstate-compose/api/restricted_current.txt b/savedstate/savedstate-compose/api/restricted_current.txt
index de12ce4..ecb7dca 100644
--- a/savedstate/savedstate-compose/api/restricted_current.txt
+++ b/savedstate/savedstate-compose/api/restricted_current.txt
@@ -22,5 +22,29 @@
     method public static inline <reified T> androidx.savedstate.compose.serialization.serializers.MutableStateSerializer<T> MutableStateSerializer();
   }
 
+  public final class SnapshotStateListSerializer<T> implements kotlinx.serialization.KSerializer<androidx.compose.runtime.snapshots.SnapshotStateList<T>> {
+    ctor public SnapshotStateListSerializer(kotlinx.serialization.KSerializer<T> elementSerializer);
+    method public androidx.compose.runtime.snapshots.SnapshotStateList<T> deserialize(kotlinx.serialization.encoding.Decoder decoder);
+    method public kotlinx.serialization.descriptors.SerialDescriptor getDescriptor();
+    method public void serialize(kotlinx.serialization.encoding.Encoder encoder, androidx.compose.runtime.snapshots.SnapshotStateList<T> value);
+    property public kotlinx.serialization.descriptors.SerialDescriptor descriptor;
+  }
+
+  public final class SnapshotStateListSerializerKt {
+    method public static inline <reified T> androidx.savedstate.compose.serialization.serializers.SnapshotStateListSerializer<T> SnapshotStateListSerializer();
+  }
+
+  public final class SnapshotStateMapSerializer<K, V> implements kotlinx.serialization.KSerializer<androidx.compose.runtime.snapshots.SnapshotStateMap<K,V>> {
+    ctor public SnapshotStateMapSerializer(kotlinx.serialization.KSerializer<K> keySerializer, kotlinx.serialization.KSerializer<V> valueSerializer);
+    method public androidx.compose.runtime.snapshots.SnapshotStateMap<K,V> deserialize(kotlinx.serialization.encoding.Decoder decoder);
+    method public kotlinx.serialization.descriptors.SerialDescriptor getDescriptor();
+    method public void serialize(kotlinx.serialization.encoding.Encoder encoder, androidx.compose.runtime.snapshots.SnapshotStateMap<K,V> value);
+    property public kotlinx.serialization.descriptors.SerialDescriptor descriptor;
+  }
+
+  public final class SnapshotStateMapSerializerKt {
+    method public static inline <reified K, reified V> androidx.savedstate.compose.serialization.serializers.SnapshotStateMapSerializer<K,V> SnapshotStateMapSerializer();
+  }
+
 }
 
diff --git a/savedstate/savedstate-compose/bcv/native/current.txt b/savedstate/savedstate-compose/bcv/native/current.txt
index 82a8223..508dbb3 100644
--- a/savedstate/savedstate-compose/bcv/native/current.txt
+++ b/savedstate/savedstate-compose/bcv/native/current.txt
@@ -6,6 +6,16 @@
 // - Show declarations: true
 
 // Library unique name: <androidx.savedstate:savedstate-compose>
+final class <#A: kotlin/Any?, #B: kotlin/Any?> androidx.savedstate.compose.serialization.serializers/SnapshotStateMapSerializer : kotlinx.serialization/KSerializer<androidx.compose.runtime.snapshots/SnapshotStateMap<#A, #B>> { // androidx.savedstate.compose.serialization.serializers/SnapshotStateMapSerializer|null[0]
+    constructor <init>(kotlinx.serialization/KSerializer<#A>, kotlinx.serialization/KSerializer<#B>) // androidx.savedstate.compose.serialization.serializers/SnapshotStateMapSerializer.<init>|<init>(kotlinx.serialization.KSerializer<1:0>;kotlinx.serialization.KSerializer<1:1>){}[0]
+
+    final val descriptor // androidx.savedstate.compose.serialization.serializers/SnapshotStateMapSerializer.descriptor|{}descriptor[0]
+        final fun <get-descriptor>(): kotlinx.serialization.descriptors/SerialDescriptor // androidx.savedstate.compose.serialization.serializers/SnapshotStateMapSerializer.descriptor.<get-descriptor>|<get-descriptor>(){}[0]
+
+    final fun deserialize(kotlinx.serialization.encoding/Decoder): androidx.compose.runtime.snapshots/SnapshotStateMap<#A, #B> // androidx.savedstate.compose.serialization.serializers/SnapshotStateMapSerializer.deserialize|deserialize(kotlinx.serialization.encoding.Decoder){}[0]
+    final fun serialize(kotlinx.serialization.encoding/Encoder, androidx.compose.runtime.snapshots/SnapshotStateMap<#A, #B>) // androidx.savedstate.compose.serialization.serializers/SnapshotStateMapSerializer.serialize|serialize(kotlinx.serialization.encoding.Encoder;androidx.compose.runtime.snapshots.SnapshotStateMap<1:0,1:1>){}[0]
+}
+
 final class <#A: kotlin/Any?> androidx.savedstate.compose.serialization.serializers/MutableStateSerializer : kotlinx.serialization/KSerializer<androidx.compose.runtime/MutableState<#A>> { // androidx.savedstate.compose.serialization.serializers/MutableStateSerializer|null[0]
     constructor <init>(kotlinx.serialization/KSerializer<#A>) // androidx.savedstate.compose.serialization.serializers/MutableStateSerializer.<init>|<init>(kotlinx.serialization.KSerializer<1:0>){}[0]
 
@@ -16,7 +26,19 @@
     final fun serialize(kotlinx.serialization.encoding/Encoder, androidx.compose.runtime/MutableState<#A>) // androidx.savedstate.compose.serialization.serializers/MutableStateSerializer.serialize|serialize(kotlinx.serialization.encoding.Encoder;androidx.compose.runtime.MutableState<1:0>){}[0]
 }
 
+final class <#A: kotlin/Any?> androidx.savedstate.compose.serialization.serializers/SnapshotStateListSerializer : kotlinx.serialization/KSerializer<androidx.compose.runtime.snapshots/SnapshotStateList<#A>> { // androidx.savedstate.compose.serialization.serializers/SnapshotStateListSerializer|null[0]
+    constructor <init>(kotlinx.serialization/KSerializer<#A>) // androidx.savedstate.compose.serialization.serializers/SnapshotStateListSerializer.<init>|<init>(kotlinx.serialization.KSerializer<1:0>){}[0]
+
+    final val descriptor // androidx.savedstate.compose.serialization.serializers/SnapshotStateListSerializer.descriptor|{}descriptor[0]
+        final fun <get-descriptor>(): kotlinx.serialization.descriptors/SerialDescriptor // androidx.savedstate.compose.serialization.serializers/SnapshotStateListSerializer.descriptor.<get-descriptor>|<get-descriptor>(){}[0]
+
+    final fun deserialize(kotlinx.serialization.encoding/Decoder): androidx.compose.runtime.snapshots/SnapshotStateList<#A> // androidx.savedstate.compose.serialization.serializers/SnapshotStateListSerializer.deserialize|deserialize(kotlinx.serialization.encoding.Decoder){}[0]
+    final fun serialize(kotlinx.serialization.encoding/Encoder, androidx.compose.runtime.snapshots/SnapshotStateList<#A>) // androidx.savedstate.compose.serialization.serializers/SnapshotStateListSerializer.serialize|serialize(kotlinx.serialization.encoding.Encoder;androidx.compose.runtime.snapshots.SnapshotStateList<1:0>){}[0]
+}
+
 final val androidx.savedstate.compose/LocalSavedStateRegistryOwner // androidx.savedstate.compose/LocalSavedStateRegistryOwner|{}LocalSavedStateRegistryOwner[0]
     final fun <get-LocalSavedStateRegistryOwner>(): androidx.compose.runtime/ProvidableCompositionLocal<androidx.savedstate/SavedStateRegistryOwner> // androidx.savedstate.compose/LocalSavedStateRegistryOwner.<get-LocalSavedStateRegistryOwner>|<get-LocalSavedStateRegistryOwner>(){}[0]
 
+final inline fun <#A: reified kotlin/Any?, #B: reified kotlin/Any?> androidx.savedstate.compose.serialization.serializers/SnapshotStateMapSerializer(): androidx.savedstate.compose.serialization.serializers/SnapshotStateMapSerializer<#A, #B> // androidx.savedstate.compose.serialization.serializers/SnapshotStateMapSerializer|SnapshotStateMapSerializer(){0§<kotlin.Any?>;1§<kotlin.Any?>}[0]
 final inline fun <#A: reified kotlin/Any?> androidx.savedstate.compose.serialization.serializers/MutableStateSerializer(): androidx.savedstate.compose.serialization.serializers/MutableStateSerializer<#A> // androidx.savedstate.compose.serialization.serializers/MutableStateSerializer|MutableStateSerializer(){0§<kotlin.Any?>}[0]
+final inline fun <#A: reified kotlin/Any?> androidx.savedstate.compose.serialization.serializers/SnapshotStateListSerializer(): androidx.savedstate.compose.serialization.serializers/SnapshotStateListSerializer<#A> // androidx.savedstate.compose.serialization.serializers/SnapshotStateListSerializer|SnapshotStateListSerializer(){0§<kotlin.Any?>}[0]
diff --git a/savedstate/savedstate-compose/src/androidInstrumentedTest/kotlin/androidx/savedstate/compose/serialization/serializers/SnapshotStateListSerializerTest.kt b/savedstate/savedstate-compose/src/androidInstrumentedTest/kotlin/androidx/savedstate/compose/serialization/serializers/SnapshotStateListSerializerTest.kt
new file mode 100644
index 0000000..133aa87
--- /dev/null
+++ b/savedstate/savedstate-compose/src/androidInstrumentedTest/kotlin/androidx/savedstate/compose/serialization/serializers/SnapshotStateListSerializerTest.kt
@@ -0,0 +1,95 @@
+/*
+ * Copyright 2025 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.savedstate.compose.serialization.serializers
+
+import androidx.compose.runtime.mutableStateListOf
+import androidx.compose.runtime.snapshots.SnapshotStateList
+import androidx.kruth.assertThat
+import androidx.savedstate.serialization.decodeFromSavedState
+import androidx.savedstate.serialization.encodeToSavedState
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import kotlin.test.Test
+import kotlinx.serialization.Serializable
+import kotlinx.serialization.serializer
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+internal class SnapshotStateListSerializerTest {
+
+    @Test
+    fun encodeDecode_serializable_withElementSerializer() {
+        doTest(mutableStateListOf(Data(1), Data(2)), SnapshotStateListSerializer(serializer()))
+    }
+
+    @Test
+    fun encodeDecode_serializable() {
+        doTest(mutableStateListOf(Data(1), Data(2)))
+    }
+
+    @Test
+    fun encodeDecode_boolean() {
+        doTest(mutableStateListOf(true, false))
+    }
+
+    @Test
+    fun encodeDecode_short() {
+        doTest(mutableStateListOf(123.toShort(), 456.toShort()))
+    }
+
+    @Test
+    fun encodeDecode_int() {
+        doTest(mutableStateListOf(123, 456))
+    }
+
+    @Test
+    fun encodeDecode_long() {
+        doTest(mutableStateListOf(123L, 456L))
+    }
+
+    @Test
+    fun encodeDecode_float() {
+        doTest(mutableStateListOf(3.14F, 2.71F))
+    }
+
+    @Test
+    fun encodeDecode_double() {
+        doTest(mutableStateListOf(3.14, 2.71))
+    }
+
+    @Test
+    fun encodeDecode_char() {
+        doTest(mutableStateListOf('c', 'd'))
+    }
+
+    @Test
+    fun encodeDecode_strings() {
+        doTest(mutableStateListOf("foo", "bar"))
+    }
+
+    private inline fun <reified T : Any> doTest(
+        original: SnapshotStateList<T>,
+        serializer: SnapshotStateListSerializer<T> = SnapshotStateListSerializer(),
+    ) {
+        val serialized = encodeToSavedState(serializer, original)
+        val deserialized = decodeFromSavedState(serializer, serialized)
+        assertThat(original.toList()).isEqualTo(deserialized.toList())
+    }
+
+    @Serializable private data class Data(val value: Int)
+}
diff --git a/savedstate/savedstate-compose/src/androidInstrumentedTest/kotlin/androidx/savedstate/compose/serialization/serializers/SnapshotStateMapSerializerTest.kt b/savedstate/savedstate-compose/src/androidInstrumentedTest/kotlin/androidx/savedstate/compose/serialization/serializers/SnapshotStateMapSerializerTest.kt
new file mode 100644
index 0000000..3766668
--- /dev/null
+++ b/savedstate/savedstate-compose/src/androidInstrumentedTest/kotlin/androidx/savedstate/compose/serialization/serializers/SnapshotStateMapSerializerTest.kt
@@ -0,0 +1,99 @@
+/*
+ * Copyright 2025 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.savedstate.compose.serialization.serializers
+
+import androidx.compose.runtime.mutableStateMapOf
+import androidx.compose.runtime.snapshots.SnapshotStateMap
+import androidx.kruth.assertThat
+import androidx.savedstate.serialization.decodeFromSavedState
+import androidx.savedstate.serialization.encodeToSavedState
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import kotlin.test.Test
+import kotlinx.serialization.Serializable
+import kotlinx.serialization.serializer
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+internal class SnapshotStateMapSerializerTest {
+
+    @Test
+    fun encodeDecode_serializable_withElementSerializer() {
+        val original = mutableStateMapOf(Data(11) to Data(21), Data(12) to Data(22))
+        val serializer = SnapshotStateMapSerializer<Data, Data>(serializer(), serializer())
+        doTest(original, serializer)
+    }
+
+    @Test
+    fun encodeDecode_serializable() {
+        val original = mutableStateMapOf(Data(11) to Data(21), Data(12) to Data(22))
+        val serializer = SnapshotStateMapSerializer<Data, Data>()
+        doTest(original, serializer)
+    }
+
+    @Test
+    fun encodeDecode_boolean() {
+        doTest(mutableStateMapOf(true to false, false to false))
+    }
+
+    @Test
+    fun encodeDecode_short() {
+        doTest(mutableStateMapOf(123.toShort() to 456.toShort(), 789.toShort() to 1011.toShort()))
+    }
+
+    @Test
+    fun encodeDecode_int() {
+        doTest(mutableStateMapOf(123 to 456, 789 to 1011))
+    }
+
+    @Test
+    fun encodeDecode_long() {
+        doTest(mutableStateMapOf(123L to 456L, 789L to 1011L))
+    }
+
+    @Test
+    fun encodeDecode_float() {
+        doTest(mutableStateMapOf(3.14F to 2.71F, 1.0F to 2.0F))
+    }
+
+    @Test
+    fun encodeDecode_double() {
+        doTest(mutableStateMapOf(3.14 to 2.71, 1.0 to 2.0))
+    }
+
+    @Test
+    fun encodeDecode_char() {
+        doTest(mutableStateMapOf('c' to 'd', 'e' to 'f'))
+    }
+
+    @Test
+    fun encodeDecode_strings() {
+        doTest(mutableStateMapOf("foo" to "bar", "baz" to "qux"))
+    }
+
+    private inline fun <reified K, reified V> doTest(
+        original: SnapshotStateMap<K, V>,
+        serializer: SnapshotStateMapSerializer<K, V> = SnapshotStateMapSerializer<K, V>()
+    ) {
+        val serialized = encodeToSavedState(serializer, original)
+        val deserialized = decodeFromSavedState(serializer, serialized)
+        assertThat(original.toMap()).isEqualTo(deserialized.toMap())
+    }
+
+    @Serializable private data class Data(val value: Int)
+}
diff --git a/savedstate/savedstate-compose/src/commonMain/kotlin/androidx/savedstate/compose/serialization/serializers/MutableStateSerializer.kt b/savedstate/savedstate-compose/src/commonMain/kotlin/androidx/savedstate/compose/serialization/serializers/MutableStateSerializer.kt
index 40892c3..c9badd7 100644
--- a/savedstate/savedstate-compose/src/commonMain/kotlin/androidx/savedstate/compose/serialization/serializers/MutableStateSerializer.kt
+++ b/savedstate/savedstate-compose/src/commonMain/kotlin/androidx/savedstate/compose/serialization/serializers/MutableStateSerializer.kt
@@ -41,7 +41,6 @@
  * @return A [MutableStateSerializer] for handling [MutableState] containing a [Serializable] type
  *   [T].
  */
-@Suppress("FunctionName")
 public inline fun <reified T> MutableStateSerializer(): MutableStateSerializer<T> {
     return MutableStateSerializer(serializer())
 }
@@ -61,22 +60,21 @@
 ) : KSerializer<MutableState<T>> {
 
     @OptIn(ExperimentalSerializationApi::class)
-    override val descriptor: SerialDescriptor by lazy {
-        val structureKind = valueSerializer.descriptor.kind
-        if (structureKind is PrimitiveKind) {
-            PrimitiveSerialDescriptor(SERIAL_NAME, structureKind)
+    override val descriptor: SerialDescriptor = run {
+        val serialName = "androidx.compose.runtime.MutableState"
+        val kind = valueSerializer.descriptor.kind
+        if (kind is PrimitiveKind) {
+            PrimitiveSerialDescriptor(serialName, kind)
         } else {
-            SerialDescriptor(SERIAL_NAME, valueSerializer.descriptor)
+            SerialDescriptor(serialName, valueSerializer.descriptor)
         }
     }
 
     override fun serialize(encoder: Encoder, value: MutableState<T>) {
-        valueSerializer.serialize(encoder, value.value)
+        encoder.encodeSerializableValue(valueSerializer, value.value)
     }
 
     override fun deserialize(decoder: Decoder): MutableState<T> {
-        return mutableStateOf(valueSerializer.deserialize(decoder))
+        return mutableStateOf(decoder.decodeSerializableValue(valueSerializer))
     }
 }
-
-private const val SERIAL_NAME = "androidx.compose.runtime.MutableState"
diff --git a/savedstate/savedstate-compose/src/commonMain/kotlin/androidx/savedstate/compose/serialization/serializers/SnapshotStateListSerializer.kt b/savedstate/savedstate-compose/src/commonMain/kotlin/androidx/savedstate/compose/serialization/serializers/SnapshotStateListSerializer.kt
new file mode 100644
index 0000000..44a1c1f
--- /dev/null
+++ b/savedstate/savedstate-compose/src/commonMain/kotlin/androidx/savedstate/compose/serialization/serializers/SnapshotStateListSerializer.kt
@@ -0,0 +1,70 @@
+/*
+ * Copyright 2025 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.savedstate.compose.serialization.serializers
+
+import androidx.compose.runtime.snapshots.SnapshotStateList
+import kotlinx.serialization.ExperimentalSerializationApi
+import kotlinx.serialization.KSerializer
+import kotlinx.serialization.builtins.ListSerializer
+import kotlinx.serialization.descriptors.SerialDescriptor
+import kotlinx.serialization.encoding.Decoder
+import kotlinx.serialization.encoding.Encoder
+import kotlinx.serialization.serializer
+
+/**
+ * Creates a [KSerializer] for a [SnapshotStateList] containing serializable elements of type [T].
+ *
+ * This inline function automatically infers the element type [T] and retrieves the corresponding
+ * [KSerializer] for serializing and deserializing [SnapshotStateList] instances.
+ *
+ * @param T The type of elements stored in the [SnapshotStateList].
+ * @return A [SnapshotStateListSerializer] for handling serialization and deserialization of a
+ *   [SnapshotStateList] containing elements of type [T].
+ */
+public inline fun <reified T> SnapshotStateListSerializer(): SnapshotStateListSerializer<T> {
+    return SnapshotStateListSerializer(serializer())
+}
+
+/**
+ * A [KSerializer] for [SnapshotStateList].
+ *
+ * This serializer wraps a [KSerializer] for the element type [T], enabling serialization and
+ * deserialization of [SnapshotStateList] instances. The serialization of individual elements is
+ * delegated to the provided [elementSerializer].
+ *
+ * @param T The type of elements stored in the [SnapshotStateList].
+ * @param elementSerializer The [KSerializer] used to serialize and deserialize individual elements.
+ */
+public class SnapshotStateListSerializer<T>(
+    private val elementSerializer: KSerializer<T>,
+) : KSerializer<SnapshotStateList<T>> {
+
+    private val base = ListSerializer(elementSerializer)
+
+    @OptIn(ExperimentalSerializationApi::class)
+    override val descriptor: SerialDescriptor =
+        SerialDescriptor("androidx.compose.runtime.SnapshotStateList", base.descriptor)
+
+    override fun serialize(encoder: Encoder, value: SnapshotStateList<T>) {
+        encoder.encodeSerializableValue(base, value)
+    }
+
+    override fun deserialize(decoder: Decoder): SnapshotStateList<T> {
+        val deserialized = decoder.decodeSerializableValue(base)
+        return SnapshotStateList<T>().apply { addAll(deserialized.toList()) }
+    }
+}
diff --git a/savedstate/savedstate-compose/src/commonMain/kotlin/androidx/savedstate/compose/serialization/serializers/SnapshotStateMapSerializer.kt b/savedstate/savedstate-compose/src/commonMain/kotlin/androidx/savedstate/compose/serialization/serializers/SnapshotStateMapSerializer.kt
new file mode 100644
index 0000000..795b2f8
--- /dev/null
+++ b/savedstate/savedstate-compose/src/commonMain/kotlin/androidx/savedstate/compose/serialization/serializers/SnapshotStateMapSerializer.kt
@@ -0,0 +1,76 @@
+/*
+ * Copyright 2025 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.savedstate.compose.serialization.serializers
+
+import androidx.compose.runtime.snapshots.SnapshotStateMap
+import kotlinx.serialization.ExperimentalSerializationApi
+import kotlinx.serialization.KSerializer
+import kotlinx.serialization.builtins.MapSerializer
+import kotlinx.serialization.descriptors.SerialDescriptor
+import kotlinx.serialization.encoding.Decoder
+import kotlinx.serialization.encoding.Encoder
+import kotlinx.serialization.serializer
+
+/**
+ * Creates a [KSerializer] for a [SnapshotStateMap] containing serializable keys of type [K] and
+ * serializable values of type [V].
+ *
+ * This inline function automatically infers the key type [K] and value type [V], and retrieves the
+ * corresponding [KSerializer] for serializing and deserializing [SnapshotStateMap] instances.
+ *
+ * @param K The type of keys stored in the [SnapshotStateMap].
+ * @param V The type of values stored in the [SnapshotStateMap].
+ * @return A [SnapshotStateMapSerializer] for handling serialization and deserialization of a
+ *   [SnapshotStateMap] containing keys of type [K] and values of type [V].
+ */
+public inline fun <reified K, reified V> SnapshotStateMapSerializer():
+    SnapshotStateMapSerializer<K, V> {
+    return SnapshotStateMapSerializer(serializer(), serializer())
+}
+
+/**
+ * A [KSerializer] for [SnapshotStateMap].
+ *
+ * This serializer wraps [KSerializer] instances for the key type [K] and value type [V], enabling
+ * serialization and deserialization of [SnapshotStateMap] instances. The serialization of
+ * individual keys and values is delegated to the provided [keySerializer] and [valueSerializer].
+ *
+ * @param K The type of keys stored in the [SnapshotStateMap].
+ * @param V The type of values stored in the [SnapshotStateMap].
+ * @param keySerializer The [KSerializer] used to serialize and deserialize individual keys.
+ * @param valueSerializer The [KSerializer] used to serialize and deserialize individual values.
+ */
+public class SnapshotStateMapSerializer<K, V>(
+    keySerializer: KSerializer<K>,
+    valueSerializer: KSerializer<V>,
+) : KSerializer<SnapshotStateMap<K, V>> {
+
+    private val base = MapSerializer(keySerializer, valueSerializer)
+
+    @OptIn(ExperimentalSerializationApi::class)
+    override val descriptor: SerialDescriptor =
+        SerialDescriptor("androidx.compose.runtime.SnapshotStateMap", base.descriptor)
+
+    override fun serialize(encoder: Encoder, value: SnapshotStateMap<K, V>) {
+        encoder.encodeSerializableValue(base, value)
+    }
+
+    override fun deserialize(decoder: Decoder): SnapshotStateMap<K, V> {
+        val deserialized = decoder.decodeSerializableValue(base)
+        return SnapshotStateMap<K, V>().apply { putAll(deserialized) }
+    }
+}
diff --git a/savedstate/savedstate/src/commonMain/kotlin/androidx/savedstate/serialization/serializers/MutableStateFlowSerializer.kt b/savedstate/savedstate/src/commonMain/kotlin/androidx/savedstate/serialization/serializers/MutableStateFlowSerializer.kt
index e9be3ba..b6f12e6 100644
--- a/savedstate/savedstate/src/commonMain/kotlin/androidx/savedstate/serialization/serializers/MutableStateFlowSerializer.kt
+++ b/savedstate/savedstate/src/commonMain/kotlin/androidx/savedstate/serialization/serializers/MutableStateFlowSerializer.kt
@@ -58,23 +58,23 @@
 public class MutableStateFlowSerializer<T>(
     private val valueSerializer: KSerializer<T>,
 ) : KSerializer<MutableStateFlow<T>> {
+
     @OptIn(ExperimentalSerializationApi::class)
-    override val descriptor: SerialDescriptor by lazy {
-        val structureKind = valueSerializer.descriptor.kind
-        if (structureKind is PrimitiveKind) {
-            PrimitiveSerialDescriptor(SERIAL_NAME, structureKind)
+    override val descriptor: SerialDescriptor = run {
+        val serialName = "kotlinx.coroutines.flow.MutableStateFlow"
+        val kind = valueSerializer.descriptor.kind
+        if (kind is PrimitiveKind) {
+            PrimitiveSerialDescriptor(serialName, kind)
         } else {
-            SerialDescriptor(SERIAL_NAME, valueSerializer.descriptor)
+            SerialDescriptor(serialName, valueSerializer.descriptor)
         }
     }
 
     override fun serialize(encoder: Encoder, value: MutableStateFlow<T>) {
-        valueSerializer.serialize(encoder, value.value)
+        encoder.encodeSerializableValue(valueSerializer, value.value)
     }
 
     override fun deserialize(decoder: Decoder): MutableStateFlow<T> {
-        return MutableStateFlow(valueSerializer.deserialize(decoder))
+        return MutableStateFlow(decoder.decodeSerializableValue(valueSerializer))
     }
 }
-
-private const val SERIAL_NAME = "kotlinx.coroutines.flow.MutableStateFlow"
diff --git a/wear/compose/compose-material3/src/main/java/androidx/wear/compose/material3/Scaffold.kt b/wear/compose/compose-material3/src/main/java/androidx/wear/compose/material3/Scaffold.kt
index 55cd5b9..c519d3f 100644
--- a/wear/compose/compose-material3/src/main/java/androidx/wear/compose/material3/Scaffold.kt
+++ b/wear/compose/compose-material3/src/main/java/androidx/wear/compose/material3/Scaffold.kt
@@ -160,4 +160,4 @@
 
 internal val LocalScaffoldState = compositionLocalOf { ScaffoldState(appScaffoldPresent = false) }
 
-private const val IDLE_DELAY = 2500L
+private const val IDLE_DELAY = 2000L