Refresh CheckoutController state after confirmation

paymentsheet/src/main/java/com/stripe/android/checkout/CheckoutConfirmationResultHandler.kt

@@ -3,6 +3,8 @@ package com.stripe.android.checkout
 import com.stripe.android.core.injection.ViewModelScope
 import com.stripe.android.paymentelement.CheckoutSessionPreview
 import com.stripe.android.paymentelement.confirmation.ConfirmationHandler
+import com.stripe.android.paymentelement.confirmation.intent.CheckoutSessionResponseKey
+import com.stripe.android.paymentsheet.repositories.CheckoutSessionResponse
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.filterIsInstance
 import kotlinx.coroutines.flow.launchIn
@@ -24,26 +26,40 @@ internal class CheckoutConfirmationResultHandler @Inject constructor(
      * per confirmation and never restores a stale completion across process death, so this single
      * collector delivers each result exactly once.
      */
-    fun register() {
+    fun register(
+        onSucceeded: suspend (CheckoutSessionResponse?) -> Unit,
+    ) {
         confirmationHandler.state
             .filterIsInstance<ConfirmationHandler.State.Complete>()
-            .onEach { handle(it.result) }
+            .onEach { handle(it.result, onSucceeded) }
             .launchIn(viewModelScope)
     }
 
-    private fun handle(result: ConfirmationHandler.Result) {
+    private suspend fun handle(
+        result: ConfirmationHandler.Result,
+        onSucceeded: suspend (CheckoutSessionResponse?) -> Unit,
+    ) {
         val isProcessDeathResult = isAwaitingProcessDeathResult
         isAwaitingProcessDeathResult = false
 
         when (result) {
-            is ConfirmationHandler.Result.Succeeded ->
-                resultCallback.onResult(CheckoutController.Result.Completed())
+            is ConfirmationHandler.Result.Succeeded -> handleSucceeded(result, onSucceeded)
             is ConfirmationHandler.Result.Failed ->
                 resultCallback.onResult(CheckoutController.Result.Failed(result.cause))
             is ConfirmationHandler.Result.Canceled -> handleCanceled(result, isProcessDeathResult)
         }
     }
 
+    private suspend fun handleSucceeded(
+        result: ConfirmationHandler.Result.Succeeded,
+        onSucceeded: suspend (CheckoutSessionResponse?) -> Unit,
+    ) {
+        // Synchronous confirmation carries the latest response. Next-action confirmation does not,
+        // so the controller retrieves the completed Checkout Session before delivering the result.
+        onSucceeded(result.metadata[CheckoutSessionResponseKey])
+        resultCallback.onResult(CheckoutController.Result.Completed())
+    }
+
     private fun handleCanceled(
         result: ConfirmationHandler.Result.Canceled,
         isProcessDeathResult: Boolean,

paymentsheet/src/main/java/com/stripe/android/checkout/CheckoutConfirmationStateUpdater.kt

@@ -0,0 +1,43 @@
+package com.stripe.android.checkout
+
+import com.stripe.android.paymentelement.CheckoutSessionPreview
+import com.stripe.android.paymentsheet.repositories.CheckoutSessionRepository
+import com.stripe.android.paymentsheet.repositories.CheckoutSessionResponse
+import javax.inject.Inject
+
+@OptIn(CheckoutSessionPreview::class)
+internal class CheckoutConfirmationStateUpdater internal constructor(
+    private val stateHolder: CheckoutControllerStateHolder,
+    private val fetchResponse: suspend (sessionId: String, adaptivePricingAllowed: Boolean) ->
+        Result<CheckoutSessionResponse>,
+    private val reloadState: suspend (CheckoutControllerState) -> Unit,
+) {
+    @Inject
+    constructor(
+        stateHolder: CheckoutControllerStateHolder,
+        checkoutSessionRepository: CheckoutSessionRepository,
+        checkoutStateLoader: CheckoutStateLoader,
+    ) : this(
+        stateHolder = stateHolder,
+        fetchResponse = checkoutSessionRepository::init,
+        reloadState = checkoutStateLoader::reload,
+    )
+
+    suspend fun update(checkoutSessionResponse: CheckoutSessionResponse?) {
+        // The controller invokes this inside runSerialized, so the snapshot remains current across
+        // the fetch and reload suspensions and cannot overwrite a concurrently committed mutation.
+        val state = stateHolder.state ?: return
+        val response = checkoutSessionResponse ?: runCatching {
+            fetchResponse(
+                state.checkoutSessionResponse.id,
+                state.configuration.adaptivePricingAllowed,
+            ).getOrThrow()
+        }.getOrNull() ?: return
+
+        // Reloading keeps every field derived from the response in sync. A refresh failure must not
+        // turn a successfully completed payment into a failed result or prevent its callback.
+        runCatching {
+            reloadState(state.copy(checkoutSessionResponse = response))
+        }
+    }
+}

paymentsheet/src/main/java/com/stripe/android/checkout/CheckoutController.kt

@@ -57,6 +57,7 @@ private val SERVER_UPDATE_TIMEOUT_MS = 20.seconds.inWholeMilliseconds
 class CheckoutController @Inject internal constructor(
     @ViewModelScope private val viewModelScope: CoroutineScope,
     confirmationResultHandler: CheckoutConfirmationResultHandler,
+    private val confirmationStateUpdater: CheckoutConfirmationStateUpdater,
     private val checkoutSessionRepository: CheckoutSessionRepository,
     private val checkoutStateLoader: CheckoutStateLoader,
     private val stateHolder: CheckoutControllerStateHolder,
@@ -70,7 +71,7 @@ class CheckoutController @Inject internal constructor(
     private val _isUpdating = MutableStateFlow(false)
 
     init {
-        confirmationResultHandler.register()
+        confirmationResultHandler.register(::handleConfirmationSucceeded)
     }
 
     /**
@@ -297,6 +298,15 @@ class CheckoutController @Inject internal constructor(
         IllegalStateException("Cannot mutate checkout session while a payment flow is presented.")
     )
 
+    internal suspend fun handleConfirmationSucceeded(
+        checkoutSessionResponse: CheckoutSessionResponse?,
+    ) {
+        runSerialized {
+            confirmationStateUpdater.update(checkoutSessionResponse)
+            kotlin.Result.success(Unit)
+        }
+    }
+
     /**
      * Serializes [block] behind [mutex] so configuration and mutations run in sequence, and toggles
      * [isUpdating] while any serialized work is in flight (tracked via [pendingMutations] so

paymentsheet/src/test/java/com/stripe/android/checkout/CheckoutConfirmationResultHandlerTest.kt

@@ -9,6 +9,10 @@ import com.stripe.android.paymentelement.CheckoutSessionPreview
 import com.stripe.android.paymentelement.confirmation.ConfirmationHandler
 import com.stripe.android.paymentelement.confirmation.FakeConfirmationHandler
 import com.stripe.android.paymentelement.confirmation.FakeConfirmationOption
+import com.stripe.android.paymentelement.confirmation.MutableConfirmationMetadata
+import com.stripe.android.paymentelement.confirmation.intent.CheckoutSessionResponseKey
+import com.stripe.android.paymentsheet.repositories.CheckoutSessionResponse
+import com.stripe.android.paymentsheet.repositories.CheckoutSessionResponseFactory
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.TestCoroutineScheduler
 import kotlinx.coroutines.test.runTest
@@ -18,30 +22,73 @@ import org.junit.Test
 internal class CheckoutConfirmationResultHandlerTest {
 
     @Test
-    fun `register handles current complete result`() {
-        runScenario(initialResult = succeeded()) {
-            assertThat(callbackCalls.awaitItem()).isInstanceOf<CheckoutController.Result.Completed>()
+    fun `register handles current complete response before invoking callback`() {
+        val response = response(status = CheckoutSessionResponse.Status.COMPLETE)
+
+        runScenario(initialResult = succeeded(response)) {
+            val callbackCall = callbackCalls.awaitItem()
+            assertThat(callbackCall.result).isInstanceOf<CheckoutController.Result.Completed>()
+            assertThat(callbackCall.succeededCallCount).isEqualTo(1)
+            assertThat(callbackCall.succeededResponse).isEqualTo(response)
         }
     }
 
     @Test
-    fun `succeeded result invokes Completed`() = runScenario {
-        emit(succeeded())
+    fun `succeeded with open response passes it to success handler and invokes Completed`() = runScenario {
+        val response = response(status = CheckoutSessionResponse.Status.OPEN)
 
-        assertThat(callbackCalls.awaitItem()).isInstanceOf<CheckoutController.Result.Completed>()
+        emit(succeeded(response))
+
+        val callbackCall = callbackCalls.awaitItem()
+        assertThat(callbackCall.result).isInstanceOf<CheckoutController.Result.Completed>()
+        assertThat(callbackCall.succeededResponse).isEqualTo(response)
+    }
+
+    @Test
+    fun `succeeded with unknown response passes it to success handler and invokes Completed`() = runScenario {
+        val response = response(status = CheckoutSessionResponse.Status.UNKNOWN)
+
+        emit(succeeded(response))
+
+        val callbackCall = callbackCalls.awaitItem()
+        assertThat(callbackCall.result).isInstanceOf<CheckoutController.Result.Completed>()
+        assertThat(callbackCall.succeededResponse).isEqualTo(response)
+    }
+
+    @Test
+    fun `succeeded with expired response passes it to success handler and invokes Completed`() = runScenario {
+        val response = response(status = CheckoutSessionResponse.Status.EXPIRED)
+
+        emit(succeeded(response))
+
+        val callbackCall = callbackCalls.awaitItem()
+        assertThat(callbackCall.result).isInstanceOf<CheckoutController.Result.Completed>()
+        assertThat(callbackCall.succeededResponse).isEqualTo(response)
+    }
+
+    @Test
+    fun `succeeded result without response invokes success handler before Completed`() = runScenario {
+        emit(ConfirmationHandler.Result.Succeeded(PaymentIntentFixtures.PI_SUCCEEDED))
+
+        val callbackCall = callbackCalls.awaitItem()
+        assertThat(callbackCall.result).isInstanceOf<CheckoutController.Result.Completed>()
+        assertThat(callbackCall.succeededCallCount).isEqualTo(1)
+        assertThat(callbackCall.succeededResponse).isNull()
     }
 
     @Test
     fun `each confirmation completion delivers its own callback`() = runScenario {
-        emit(succeeded())
-        assertThat(callbackCalls.awaitItem()).isInstanceOf<CheckoutController.Result.Completed>()
+        val response = response(status = CheckoutSessionResponse.Status.COMPLETE)
+
+        emit(succeeded(response))
+        assertThat(callbackCalls.awaitItem().result).isInstanceOf<CheckoutController.Result.Completed>()
 
         // A second confirmation returns to Confirming before completing again. The interleaved
         // Confirming breaks the state flow's de-duplication, so an identical result still delivers.
         emitConfirming()
-        emit(succeeded())
+        emit(succeeded(response))
 
-        assertThat(callbackCalls.awaitItem()).isInstanceOf<CheckoutController.Result.Completed>()
+        assertThat(callbackCalls.awaitItem().result).isInstanceOf<CheckoutController.Result.Completed>()
     }
 
     @Test
@@ -56,7 +103,7 @@ internal class CheckoutConfirmationResultHandlerTest {
             )
         )
 
-        val result = callbackCalls.awaitItem()
+        val result = callbackCalls.awaitItem().result
         assertThat(result).isInstanceOf<CheckoutController.Result.Failed>()
         assertThat((result as CheckoutController.Result.Failed).error).isEqualTo(cause)
     }
@@ -69,7 +116,7 @@ internal class CheckoutConfirmationResultHandlerTest {
             )
         )
 
-        assertThat(callbackCalls.awaitItem()).isInstanceOf<CheckoutController.Result.Canceled>()
+        assertThat(callbackCalls.awaitItem().result).isInstanceOf<CheckoutController.Result.Canceled>()
     }
 
     @Test
@@ -104,15 +151,17 @@ internal class CheckoutConfirmationResultHandlerTest {
             )
         )
 
-        assertThat(callbackCalls.awaitItem()).isInstanceOf<CheckoutController.Result.Canceled>()
+        assertThat(callbackCalls.awaitItem().result).isInstanceOf<CheckoutController.Result.Canceled>()
     }
 
     private fun runScenario(
         initialResult: ConfirmationHandler.Result? = null,
         hasReloadedFromProcessDeath: Boolean = false,
         block: suspend Scenario.() -> Unit,
     ) = runTest {
-        val callbackCalls = Turbine<CheckoutController.Result>()
+        val callbackCalls = Turbine<CallbackCall>()
+        var succeededCallCount = 0
+        var succeededResponse: CheckoutSessionResponse? = null
         val initialConfirmationState: ConfirmationHandler.State =
             initialResult?.let(ConfirmationHandler.State::Complete)
                 ?: ConfirmationHandler.State.Idle
@@ -122,10 +171,21 @@ internal class CheckoutConfirmationResultHandlerTest {
         )
         val handler = CheckoutConfirmationResultHandler(
             confirmationHandler = confirmationHandler,
-            resultCallback = CheckoutController.ResultCallback(callbackCalls::add),
+            resultCallback = CheckoutController.ResultCallback { result ->
+                callbackCalls.add(
+                    CallbackCall(
+                        result = result,
+                        succeededCallCount = succeededCallCount,
+                        succeededResponse = succeededResponse,
+                    )
+                )
+            },
             viewModelScope = backgroundScope,
         )
-        handler.register()
+        handler.register { response ->
+            succeededCallCount += 1
+            succeededResponse = response
+        }
         testScheduler.runCurrent()
 
         Scenario(
@@ -138,13 +198,29 @@ internal class CheckoutConfirmationResultHandlerTest {
         callbackCalls.ensureAllEventsConsumed()
     }
 
-    private fun succeeded(): ConfirmationHandler.Result.Succeeded {
-        return ConfirmationHandler.Result.Succeeded(PaymentIntentFixtures.PI_SUCCEEDED)
+    private fun response(
+        status: CheckoutSessionResponse.Status,
+    ): CheckoutSessionResponse {
+        return CheckoutSessionResponseFactory.create(
+            id = "cs_confirmed",
+            status = status,
+        )
+    }
+
+    private fun succeeded(
+        response: CheckoutSessionResponse,
+    ): ConfirmationHandler.Result.Succeeded {
+        return ConfirmationHandler.Result.Succeeded(
+            intent = PaymentIntentFixtures.PI_SUCCEEDED,
+            metadata = MutableConfirmationMetadata().apply {
+                set(CheckoutSessionResponseKey, response)
+            },
+        )
     }
 
     private class Scenario(
         val confirmationHandler: FakeConfirmationHandler,
-        val callbackCalls: Turbine<CheckoutController.Result>,
+        val callbackCalls: Turbine<CallbackCall>,
         val testScheduler: TestCoroutineScheduler,
     ) {
         fun emit(result: ConfirmationHandler.Result) {
@@ -158,4 +234,10 @@ internal class CheckoutConfirmationResultHandlerTest {
             testScheduler.runCurrent()
         }
     }
+
+    private data class CallbackCall(
+        val result: CheckoutController.Result,
+        val succeededCallCount: Int,
+        val succeededResponse: CheckoutSessionResponse?,
+    )
 }