Merge changes I4cc652fe,I03f43442 into main

* changes:
  PriorityNestedScrollConnection: make onStop suspendable and add onCancel
  Simplify PriorityNestedScrollConnection
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationScrimNestedScrollConnection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationScrimNestedScrollConnection.kt
index 8b9e927..e4c60e1 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationScrimNestedScrollConnection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationScrimNestedScrollConnection.kt
@@ -18,6 +18,8 @@
 
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
+import androidx.compose.ui.util.fastCoerceAtLeast
+import androidx.compose.ui.util.fastCoerceAtMost
 import com.android.compose.nestedscroll.PriorityNestedScrollConnection
 
 /**
@@ -44,7 +46,7 @@
         orientation = Orientation.Vertical,
         // scrolling up and inner content is taller than the scrim, so scrim needs to
         // expand; content can scroll once scrim is at the minScrimOffset.
-        canStartPreScroll = { offsetAvailable, offsetBeforeStart ->
+        canStartPreScroll = { offsetAvailable, offsetBeforeStart, _ ->
             offsetAvailable < 0 &&
                 offsetBeforeStart == 0f &&
                 contentHeight() > minVisibleScrimHeight() &&
@@ -52,36 +54,38 @@
         },
         // scrolling down and content is done scrolling to top. After that, the scrim
         // needs to collapse; collapse the scrim until it is at the maxScrimOffset.
-        canStartPostScroll = { offsetAvailable, _ ->
+        canStartPostScroll = { offsetAvailable, _, _ ->
             offsetAvailable > 0 && (scrimOffset() < maxScrimOffset || isCurrentGestureOverscroll())
         },
         canStartPostFling = { false },
-        canContinueScroll = {
-            val currentHeight = scrimOffset()
-            minScrimOffset() < currentHeight && currentHeight < maxScrimOffset
-        },
-        canScrollOnFling = true,
+        canStopOnPreFling = { false },
         onStart = { offsetAvailable -> onStart(offsetAvailable) },
-        onScroll = { offsetAvailable ->
+        onScroll = { offsetAvailable, _ ->
             val currentHeight = scrimOffset()
             val amountConsumed =
                 if (offsetAvailable > 0) {
                     val amountLeft = maxScrimOffset - currentHeight
-                    offsetAvailable.coerceAtMost(amountLeft)
+                    offsetAvailable.fastCoerceAtMost(amountLeft)
                 } else {
                     val amountLeft = minScrimOffset() - currentHeight
-                    offsetAvailable.coerceAtLeast(amountLeft)
+                    offsetAvailable.fastCoerceAtLeast(amountLeft)
                 }
             snapScrimOffset(currentHeight + amountConsumed)
             amountConsumed
         },
-        // Don't consume the velocity on pre/post fling
         onStop = { velocityAvailable ->
             onStop(velocityAvailable)
             if (scrimOffset() < minScrimOffset()) {
                 animateScrimOffset(minScrimOffset())
             }
-            { 0f }
+            // Don't consume the velocity on pre/post fling
+            0f
+        },
+        onCancel = {
+            onStop(0f)
+            if (scrimOffset() < minScrimOffset()) {
+                animateScrimOffset(minScrimOffset())
+            }
         },
     )
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationStackNestedScrollConnection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationStackNestedScrollConnection.kt
index a706585..edb05eb 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationStackNestedScrollConnection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationStackNestedScrollConnection.kt
@@ -28,6 +28,7 @@
 import androidx.compose.ui.platform.LocalDensity
 import androidx.compose.ui.unit.IntOffset
 import androidx.compose.ui.unit.dp
+import androidx.compose.ui.util.fastCoerceAtLeast
 import com.android.compose.nestedscroll.PriorityNestedScrollConnection
 import kotlin.math.max
 import kotlin.math.roundToInt
@@ -86,21 +87,25 @@
 ): PriorityNestedScrollConnection {
     return PriorityNestedScrollConnection(
         orientation = Orientation.Vertical,
-        canStartPreScroll = { _, _ -> false },
-        canStartPostScroll = { offsetAvailable, offsetBeforeStart ->
+        canStartPreScroll = { _, _, _ -> false },
+        canStartPostScroll = { offsetAvailable, offsetBeforeStart, _ ->
             offsetAvailable < 0f && offsetBeforeStart < 0f && !canScrollForward()
         },
         canStartPostFling = { velocityAvailable -> velocityAvailable < 0f && !canScrollForward() },
-        canContinueScroll = { stackOffset() > 0f },
-        canScrollOnFling = true,
+        canStopOnPreFling = { false },
         onStart = { offsetAvailable -> onStart(offsetAvailable) },
-        onScroll = { offsetAvailable ->
-            onScroll(offsetAvailable)
-            offsetAvailable
+        onScroll = { offsetAvailable, _ ->
+            val minOffset = 0f
+            val consumed = offsetAvailable.fastCoerceAtLeast(minOffset - stackOffset())
+            if (consumed != 0f) {
+                onScroll(consumed)
+            }
+            consumed
         },
         onStop = { velocityAvailable ->
             onStop(velocityAvailable)
-            suspend { velocityAvailable }
+            velocityAvailable
         },
+        onCancel = { onStop(0f) },
     )
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
index 085157a..7e288dd 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
@@ -27,9 +27,10 @@
 import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.animation.scene.content.state.TransitionState.HasOverscrollProperties.Companion.DistanceUnspecified
 import com.android.compose.nestedscroll.PriorityNestedScrollConnection
-import com.android.compose.nestedscroll.SuspendedValue
 import kotlin.math.absoluteValue
 
+internal typealias SuspendedValue<T> = suspend () -> T
+
 internal interface DraggableHandler {
     /**
      * Start a drag in the given [startedPosition], with the given [overSlop] and number of
@@ -612,7 +613,7 @@
 
         return PriorityNestedScrollConnection(
             orientation = orientation,
-            canStartPreScroll = { offsetAvailable, offsetBeforeStart ->
+            canStartPreScroll = { offsetAvailable, offsetBeforeStart, _ ->
                 canChangeScene =
                     if (isExternalOverscrollGesture()) false else offsetBeforeStart == 0f
 
@@ -638,7 +639,7 @@
                 isIntercepting = true
                 true
             },
-            canStartPostScroll = { offsetAvailable, offsetBeforeStart ->
+            canStartPostScroll = { offsetAvailable, offsetBeforeStart, _ ->
                 val behavior: NestedScrollBehavior =
                     when {
                         offsetAvailable > 0f -> topOrLeftBehavior
@@ -693,8 +694,7 @@
 
                 canStart
             },
-            canContinueScroll = { true },
-            canScrollOnFling = false,
+            canStopOnPreFling = { true },
             onStart = { offsetAvailable ->
                 val pointersInfo = pointersInfo()
                 dragController =
@@ -704,7 +704,7 @@
                         overSlop = if (isIntercepting) 0f else offsetAvailable,
                     )
             },
-            onScroll = { offsetAvailable ->
+            onScroll = { offsetAvailable, _ ->
                 val controller = dragController ?: error("Should be called after onStart")
 
                 // TODO(b/297842071) We should handle the overscroll or slow drag if the gesture is
@@ -713,10 +713,18 @@
             },
             onStop = { velocityAvailable ->
                 val controller = dragController ?: error("Should be called after onStart")
-
-                controller
-                    .onStop(velocity = velocityAvailable, canChangeContent = canChangeScene)
-                    .also { dragController = null }
+                try {
+                    controller
+                        .onStop(velocity = velocityAvailable, canChangeContent = canChangeScene)
+                        .invoke()
+                } finally {
+                    dragController = null
+                }
+            },
+            onCancel = {
+                val controller = dragController ?: error("Should be called after onStart")
+                controller.onStop(velocity = 0f, canChangeContent = canChangeScene)
+                dragController = null
             },
         )
     }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeAnimation.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeAnimation.kt
index 205267d..f0043e1 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeAnimation.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeAnimation.kt
@@ -27,7 +27,6 @@
 import androidx.compose.ui.unit.IntSize
 import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.animation.scene.content.state.TransitionState.HasOverscrollProperties.Companion.DistanceUnspecified
-import com.android.compose.nestedscroll.SuspendedValue
 import kotlin.math.absoluteValue
 import kotlinx.coroutines.CompletableDeferred
 
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/LargeTopAppBarNestedScrollConnection.kt b/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/LargeTopAppBarNestedScrollConnection.kt
index 4ae3235..ecf64b7 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/LargeTopAppBarNestedScrollConnection.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/LargeTopAppBarNestedScrollConnection.kt
@@ -18,6 +18,8 @@
 
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
+import androidx.compose.ui.util.fastCoerceAtLeast
+import androidx.compose.ui.util.fastCoerceAtMost
 
 /**
  * A [NestedScrollConnection] that listens for all vertical scroll events and responds in the
@@ -43,35 +45,32 @@
         orientation = Orientation.Vertical,
         // When swiping up, the LargeTopAppBar will shrink (to [minHeight]) and the content will
         // expand. Then, you can then scroll down the content.
-        canStartPreScroll = { offsetAvailable, offsetBeforeStart ->
+        canStartPreScroll = { offsetAvailable, offsetBeforeStart, _ ->
             offsetAvailable < 0 && offsetBeforeStart == 0f && height() > minHeight()
         },
         // When swiping down, the content will scroll up until it reaches the top. Then, the
         // LargeTopAppBar will expand until it reaches its [maxHeight].
-        canStartPostScroll = { offsetAvailable, _ ->
+        canStartPostScroll = { offsetAvailable, _, _ ->
             offsetAvailable > 0 && height() < maxHeight()
         },
         canStartPostFling = { false },
-        canContinueScroll = {
-            val currentHeight = height()
-            minHeight() < currentHeight && currentHeight < maxHeight()
-        },
-        canScrollOnFling = true,
+        canStopOnPreFling = { false },
         onStart = { /* do nothing */ },
-        onScroll = { offsetAvailable ->
+        onScroll = { offsetAvailable, _ ->
             val currentHeight = height()
             val amountConsumed =
                 if (offsetAvailable > 0) {
                     val amountLeft = maxHeight() - currentHeight
-                    offsetAvailable.coerceAtMost(amountLeft)
+                    offsetAvailable.fastCoerceAtMost(amountLeft)
                 } else {
                     val amountLeft = minHeight() - currentHeight
-                    offsetAvailable.coerceAtLeast(amountLeft)
+                    offsetAvailable.fastCoerceAtLeast(amountLeft)
                 }
             onHeightChanged(currentHeight + amountConsumed)
             amountConsumed
         },
         // Don't consume the velocity on pre/post fling
-        onStop = { { 0f } },
+        onStop = { 0f },
+        onCancel = { /* do nothing */ },
     )
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/PriorityNestedScrollConnection.kt b/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/PriorityNestedScrollConnection.kt
index a3641e6..636c557 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/PriorityNestedScrollConnection.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/PriorityNestedScrollConnection.kt
@@ -16,37 +16,59 @@
 
 package com.android.compose.nestedscroll
 
+import androidx.compose.animation.core.AnimationState
+import androidx.compose.animation.core.DecayAnimationSpec
+import androidx.compose.animation.core.animateDecay
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
 import androidx.compose.ui.input.nestedscroll.NestedScrollSource
 import androidx.compose.ui.unit.Velocity
 import com.android.compose.ui.util.SpaceVectorConverter
+import kotlin.math.abs
 import kotlin.math.sign
-
-internal typealias SuspendedValue<T> = suspend () -> T
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.async
+import kotlinx.coroutines.coroutineScope
 
 /**
- * This [NestedScrollConnection] waits for a child to scroll ([onPreScroll] or [onPostScroll]), and
- * then decides (via [canStartPreScroll] or [canStartPostScroll]) if it should take over scrolling.
- * If it does, it will scroll before its children, until [canContinueScroll] allows it.
+ * A [NestedScrollConnection] that intercepts scroll events in priority mode.
  *
- * Note: Call [reset] before destroying this object to make sure you always get a call to [onStop]
- * after [onStart].
+ * Priority mode allows this connection to take control over scroll events within a nested scroll
+ * hierarchy. When in priority mode, this connection consumes scroll events before its children,
+ * enabling custom scrolling behaviors like sticky headers.
  *
+ * @param orientation The orientation of the scroll.
+ * @param canStartPreScroll lambda that returns true if the connection can start consuming scroll
+ *   events in pre-scroll mode.
+ * @param canStartPostScroll lambda that returns true if the connection can start consuming scroll
+ *   events in post-scroll mode.
+ * @param canStartPostFling lambda that returns true if the connection can start consuming scroll
+ *   events in post-fling mode.
+ * @param canStopOnPreFling lambda that returns true if the connection can stop consuming scroll
+ *   events in pre-fling (i.e. as soon as the user lifts their fingers).
+ * @param onStart lambda that is called when the connection starts consuming scroll events.
+ * @param onScroll lambda that is called when the connection consumes a scroll event and returns the
+ *   consumed amount.
+ * @param onStop lambda that is called when the connection stops consuming scroll events and returns
+ *   the consumed velocity.
+ * @param onCancel lambda that is called when the connection is cancelled.
  * @sample LargeTopAppBarNestedScrollConnection
  * @sample com.android.compose.animation.scene.NestedScrollHandlerImpl.nestedScrollConnection
  */
 class PriorityNestedScrollConnection(
     orientation: Orientation,
-    private val canStartPreScroll: (offsetAvailable: Float, offsetBeforeStart: Float) -> Boolean,
-    private val canStartPostScroll: (offsetAvailable: Float, offsetBeforeStart: Float) -> Boolean,
+    private val canStartPreScroll:
+        (offsetAvailable: Float, offsetBeforeStart: Float, source: NestedScrollSource) -> Boolean,
+    private val canStartPostScroll:
+        (offsetAvailable: Float, offsetBeforeStart: Float, source: NestedScrollSource) -> Boolean,
     private val canStartPostFling: (velocityAvailable: Float) -> Boolean,
-    private val canContinueScroll: (source: NestedScrollSource) -> Boolean,
-    private val canScrollOnFling: Boolean,
+    private val canStopOnPreFling: () -> Boolean,
     private val onStart: (offsetAvailable: Float) -> Unit,
-    private val onScroll: (offsetAvailable: Float) -> Float,
-    private val onStop: (velocityAvailable: Float) -> SuspendedValue<Float>,
+    private val onScroll: (offsetAvailable: Float, source: NestedScrollSource) -> Float,
+    private val onStop: suspend (velocityAvailable: Float) -> Float,
+    private val onCancel: () -> Unit,
 ) : NestedScrollConnection, SpaceVectorConverter by SpaceVectorConverter(orientation) {
 
     /** In priority mode [onPreScroll] events are first consumed by the parent, via [onScroll]. */
@@ -54,6 +76,9 @@
 
     private var offsetScrolledBeforePriorityMode = 0f
 
+    /** This job allows us to interrupt the onStop animation */
+    private var onStopJob: Deferred<Float> = CompletableDeferred(0f)
+
     override fun onPostScroll(
         consumed: Offset,
         available: Offset,
@@ -64,62 +89,48 @@
         // the beginning or from the last fling gesture.
         val offsetBeforeStart = offsetScrolledBeforePriorityMode - availableFloat
 
-        if (
-            isPriorityMode ||
-                (source == NestedScrollSource.SideEffect && !canScrollOnFling) ||
-                !canStartPostScroll(availableFloat, offsetBeforeStart)
-        ) {
+        if (isPriorityMode || !canStartPostScroll(availableFloat, offsetBeforeStart, source)) {
             // The priority mode cannot start so we won't consume the available offset.
             return Offset.Zero
         }
 
-        return onPriorityStart(availableFloat).toOffset()
+        return start(availableFloat, source).toOffset()
     }
 
     override fun onPreScroll(available: Offset, source: NestedScrollSource): Offset {
         if (!isPriorityMode) {
-            if (source == NestedScrollSource.UserInput || canScrollOnFling) {
-                val availableFloat = available.toFloat()
-                if (canStartPreScroll(availableFloat, offsetScrolledBeforePriorityMode)) {
-                    return onPriorityStart(availableFloat).toOffset()
-                }
-                // We want to track the amount of offset consumed before entering priority mode
-                offsetScrolledBeforePriorityMode += availableFloat
+            val availableFloat = available.toFloat()
+            if (canStartPreScroll(availableFloat, offsetScrolledBeforePriorityMode, source)) {
+                return start(availableFloat, source).toOffset()
             }
-
-            return Offset.Zero
-        }
-
-        val availableFloat = available.toFloat()
-        if (!canContinueScroll(source)) {
-            // Step 3a: We have lost priority and we no longer need to intercept scroll events.
-            onPriorityStop(velocity = 0f)
-
-            // We've just reset offsetScrolledBeforePriorityMode to 0f
             // We want to track the amount of offset consumed before entering priority mode
             offsetScrolledBeforePriorityMode += availableFloat
-
             return Offset.Zero
         }
 
-        // Step 2: We have the priority and can consume the scroll events.
-        return onScroll(availableFloat).toOffset()
+        return scroll(available.toFloat(), source).toOffset()
     }
 
     override suspend fun onPreFling(available: Velocity): Velocity {
-        if (isPriorityMode && canScrollOnFling) {
-            // We don't want to consume the velocity, we prefer to continue receiving scroll events.
+        if (!isPriorityMode) {
+            resetOffsetTracker()
             return Velocity.Zero
         }
-        // Step 3b: The finger is lifted, we can stop intercepting scroll events and use the speed
-        // of the fling gesture.
-        return onPriorityStop(velocity = available.toFloat()).invoke().toVelocity()
+
+        if (canStopOnPreFling()) {
+            // Step 3b: The finger is lifted, we can stop intercepting scroll events and use the
+            // velocity of the fling gesture.
+            return stop(velocityAvailable = available.toFloat()).toVelocity()
+        }
+
+        // We don't want to consume the velocity, we prefer to continue receiving scroll events.
+        return Velocity.Zero
     }
 
     override suspend fun onPostFling(consumed: Velocity, available: Velocity): Velocity {
         val availableFloat = available.toFloat()
         if (isPriorityMode) {
-            return onPriorityStop(velocity = availableFloat).invoke().toVelocity()
+            return stop(velocityAvailable = availableFloat).toVelocity()
         }
 
         if (!canStartPostFling(availableFloat)) {
@@ -131,10 +142,14 @@
         // TODO(b/291053278): Remove canStartPostFling() and instead make it possible to define the
         // overscroll behavior on the Scene level.
         val smallOffset = availableFloat.sign
-        onPriorityStart(availableOffset = smallOffset)
+        start(
+            availableOffset = smallOffset,
+            source = NestedScrollSource.SideEffect,
+            skipScroll = true,
+        )
 
         // This is the last event of a scroll gesture.
-        return onPriorityStop(availableFloat).invoke().toVelocity()
+        return stop(availableFloat).toVelocity()
     }
 
     /**
@@ -143,36 +158,76 @@
      * TODO(b/303224944) This method should be removed.
      */
     fun reset() {
-        // Step 3c: To ensure that an onStop is always called for every onStart.
-        onPriorityStop(velocity = 0f)
+        if (isPriorityMode) {
+            // Step 3c: To ensure that an onStop (or onCancel) is always called for every onStart.
+            cancel()
+        } else {
+            resetOffsetTracker()
+        }
     }
 
-    private fun onPriorityStart(availableOffset: Float): Float {
-        if (isPriorityMode) {
-            error("This should never happen, onPriorityStart() was called when isPriorityMode")
+    private fun shouldStop(consumed: Float): Boolean {
+        return consumed == 0f
+    }
+
+    private fun start(
+        availableOffset: Float,
+        source: NestedScrollSource,
+        skipScroll: Boolean = false,
+    ): Float {
+        check(!isPriorityMode) {
+            "This should never happen, start() was called when isPriorityMode"
         }
 
         // Step 1: It's our turn! We start capturing scroll events when one of our children has an
         // available offset following a scroll event.
         isPriorityMode = true
 
+        onStopJob.cancel()
+
         // Note: onStop will be called if we cannot continue to scroll (step 3a), or the finger is
         // lifted (step 3b), or this object has been destroyed (step 3c).
         onStart(availableOffset)
 
-        return onScroll(availableOffset)
+        return if (skipScroll) 0f else scroll(availableOffset, source)
     }
 
-    private fun onPriorityStop(velocity: Float): SuspendedValue<Float> {
-        // We can restart tracking the consumed offsets from scratch.
-        offsetScrolledBeforePriorityMode = 0f
+    private fun scroll(offsetAvailable: Float, source: NestedScrollSource): Float {
+        // Step 2: We have the priority and can consume the scroll events.
+        val consumedByScroll = onScroll(offsetAvailable, source)
 
-        if (!isPriorityMode) {
-            return { 0f }
+        if (shouldStop(consumedByScroll)) {
+            // Step 3a: We have lost priority and we no longer need to intercept scroll events.
+            cancel()
+
+            // We've just reset offsetScrolledBeforePriorityMode to 0f
+            // We want to track the amount of offset consumed before entering priority mode
+            offsetScrolledBeforePriorityMode += offsetAvailable - consumedByScroll
         }
 
-        isPriorityMode = false
+        return consumedByScroll
+    }
 
-        return onStop(velocity)
+    /** Reset the tracking of consumed offsets before entering in priority mode. */
+    private fun resetOffsetTracker() {
+        offsetScrolledBeforePriorityMode = 0f
+    }
+
+    private suspend fun stop(velocityAvailable: Float): Float {
+        check(isPriorityMode) { "This should never happen, stop() was called before start()" }
+        isPriorityMode = false
+        resetOffsetTracker()
+
+        return coroutineScope {
+            onStopJob = async { onStop(velocityAvailable) }
+            onStopJob.await()
+        }
+    }
+
+    private fun cancel() {
+        check(isPriorityMode) { "This should never happen, cancel() was called before start()" }
+        isPriorityMode = false
+        resetOffsetTracker()
+        onCancel()
     }
 }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
index ecef6be..57b9423 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
@@ -39,7 +39,6 @@
 import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.animation.scene.content.state.TransitionState.Transition
 import com.android.compose.animation.scene.subjects.assertThat
-import com.android.compose.nestedscroll.SuspendedValue
 import com.android.compose.test.MonotonicClockTestScope
 import com.android.compose.test.runMonotonicClockTest
 import com.android.compose.test.transition
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt
index c8f6e6d..3df6087 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt
@@ -46,7 +46,6 @@
 import androidx.compose.ui.unit.Velocity
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.compose.modifiers.thenIf
-import com.android.compose.nestedscroll.SuspendedValue
 import com.google.common.truth.Truth.assertThat
 import kotlin.properties.Delegates
 import kotlinx.coroutines.coroutineScope
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/nestedscroll/PriorityNestedScrollConnectionTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/nestedscroll/PriorityNestedScrollConnectionTest.kt
index badc43b..1a3b86b 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/nestedscroll/PriorityNestedScrollConnectionTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/nestedscroll/PriorityNestedScrollConnectionTest.kt
@@ -34,30 +34,31 @@
     private var canStartPreScroll = false
     private var canStartPostScroll = false
     private var canStartPostFling = false
-    private var canContinueScroll = false
+    private var canStopOnPreFling = true
     private var isStarted = false
     private var lastScroll: Float? = null
-    private var returnOnScroll = 0f
+    private var consumeScroll = true
     private var lastStop: Float? = null
-    private var returnOnStop = 0f
+    private var isCancelled: Boolean = false
+    private var consumeStop = true
 
     private val scrollConnection =
         PriorityNestedScrollConnection(
             orientation = Orientation.Vertical,
-            canStartPreScroll = { _, _ -> canStartPreScroll },
-            canStartPostScroll = { _, _ -> canStartPostScroll },
+            canStartPreScroll = { _, _, _ -> canStartPreScroll },
+            canStartPostScroll = { _, _, _ -> canStartPostScroll },
             canStartPostFling = { canStartPostFling },
-            canContinueScroll = { canContinueScroll },
-            canScrollOnFling = false,
+            canStopOnPreFling = { canStopOnPreFling },
             onStart = { isStarted = true },
-            onScroll = {
-                lastScroll = it
-                returnOnScroll
+            onScroll = { offsetAvailable, _ ->
+                lastScroll = offsetAvailable
+                if (consumeScroll) offsetAvailable else 0f
             },
             onStop = {
                 lastStop = it
-                { returnOnStop }
+                if (consumeStop) it else 0f
             },
+            onCancel = { isCancelled = true },
         )
 
     @Test
@@ -85,7 +86,7 @@
         canStartPostScroll = true
         scrollConnection.onPostScroll(
             consumed = Offset.Zero,
-            available = Offset.Zero,
+            available = Offset(1f, 1f),
             source = UserInput,
         )
     }
@@ -136,45 +137,55 @@
     @Test
     fun step2_onPriorityMode_shouldContinueIfAllowed() {
         startPriorityModePostScroll()
-        canContinueScroll = true
 
-        scrollConnection.onPreScroll(available = Offset(1f, 1f), source = UserInput)
+        val scroll1 = scrollConnection.onPreScroll(available = Offset(0f, 1f), source = UserInput)
         assertThat(lastScroll).isEqualTo(1f)
+        assertThat(scroll1.y).isEqualTo(1f)
 
-        canContinueScroll = false
-        scrollConnection.onPreScroll(available = Offset(2f, 2f), source = UserInput)
-        assertThat(lastScroll).isNotEqualTo(2f)
-        assertThat(lastScroll).isEqualTo(1f)
+        consumeScroll = false
+        val scroll2 = scrollConnection.onPreScroll(available = Offset(0f, 2f), source = UserInput)
+        assertThat(lastScroll).isEqualTo(2f)
+        assertThat(scroll2.y).isEqualTo(0f)
     }
 
     @Test
-    fun step3a_onPriorityMode_shouldStopIfCannotContinue() {
+    fun step3a_onPriorityMode_shouldCancelIfCannotContinue() {
         startPriorityModePostScroll()
-        canContinueScroll = false
+        consumeScroll = false
 
-        scrollConnection.onPreScroll(available = Offset.Zero, source = UserInput)
+        scrollConnection.onPreScroll(available = Offset(0f, 1f), source = UserInput)
 
-        assertThat(lastStop).isNotNull()
+        assertThat(isCancelled).isTrue()
     }
 
     @Test
     fun step3b_onPriorityMode_shouldStopOnFling() = runTest {
         startPriorityModePostScroll()
-        canContinueScroll = true
 
         scrollConnection.onPreFling(available = Velocity.Zero)
 
-        assertThat(lastStop).isNotNull()
+        assertThat(lastStop).isEqualTo(0f)
     }
 
     @Test
-    fun step3c_onPriorityMode_shouldStopOnReset() {
+    fun ifCannotStopOnPreFling_shouldStopOnPostFling() = runTest {
         startPriorityModePostScroll()
-        canContinueScroll = true
+        canStopOnPreFling = false
+
+        scrollConnection.onPreFling(available = Velocity.Zero)
+        assertThat(lastStop).isNull()
+
+        scrollConnection.onPostFling(consumed = Velocity.Zero, available = Velocity.Zero)
+        assertThat(lastStop).isEqualTo(0f)
+    }
+
+    @Test
+    fun step3c_onPriorityMode_shouldCancelOnReset() {
+        startPriorityModePostScroll()
 
         scrollConnection.reset()
 
-        assertThat(lastStop).isNotNull()
+        assertThat(isCancelled).isTrue()
     }
 
     @Test
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationScrimNestedScrollConnectionTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationScrimNestedScrollConnectionTest.kt
index 35e4047..97fa6eb 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationScrimNestedScrollConnectionTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationScrimNestedScrollConnectionTest.kt
@@ -31,6 +31,7 @@
 @RunWith(AndroidJUnit4::class)
 class NotificationScrimNestedScrollConnectionTest : SysuiTestCase() {
     private var isStarted = false
+    private var wasStarted = false
     private var scrimOffset = 0f
     private var contentHeight = 0f
     private var isCurrentGestureOverscroll = false
@@ -46,7 +47,10 @@
             minVisibleScrimHeight = { MIN_VISIBLE_SCRIM_HEIGHT },
             isCurrentGestureOverscroll = { isCurrentGestureOverscroll },
             onStart = { isStarted = true },
-            onStop = { isStarted = false },
+            onStop = {
+                wasStarted = true
+                isStarted = false
+            },
         )
 
     @Test
@@ -180,6 +184,7 @@
             )
 
         assertThat(offsetConsumed).isEqualTo(Offset.Zero)
+        assertThat(wasStarted).isEqualTo(false)
         assertThat(isStarted).isEqualTo(false)
     }
 
@@ -196,7 +201,9 @@
             )
 
         assertThat(offsetConsumed).isEqualTo(Offset.Zero)
-        assertThat(isStarted).isEqualTo(true)
+        // Returning 0 offset will immediately stop the connection
+        assertThat(wasStarted).isEqualTo(true)
+        assertThat(isStarted).isEqualTo(false)
     }
 
     @Test