Keep reference to pending payment result

paymentsheet/src/main/java/com/stripe/android/paymentsheet/PaymentSheetViewModel.kt

@@ -55,6 +55,7 @@ import com.stripe.android.paymentsheet.state.GooglePayState
 import com.stripe.android.paymentsheet.state.PaymentSheetLoader
 import com.stripe.android.paymentsheet.state.PaymentSheetState
 import com.stripe.android.paymentsheet.state.WalletsState
+import com.stripe.android.paymentsheet.state.asPaymentSheetLoadingException
 import com.stripe.android.paymentsheet.ui.HeaderTextFactory
 import com.stripe.android.paymentsheet.ui.ModifiableEditPaymentMethodViewInteractor
 import com.stripe.android.paymentsheet.ui.PrimaryButton
@@ -180,6 +181,8 @@ internal class PaymentSheetViewModel @Inject internal constructor(
             null -> GooglePayButtonType.Pay
         }
 
+    private var pendingPaymentResult: InternalPaymentResult? = null
+
     @VisibleForTesting
     internal val googlePayLauncherConfig: GooglePayPaymentMethodLauncher.Config? =
         args.googlePayConfig?.let { config ->
@@ -318,16 +321,27 @@ internal class PaymentSheetViewModel @Inject internal constructor(
         }
 
         result.fold(
-            onSuccess = { state ->
-                handlePaymentSheetStateLoaded(state)
-            },
-            onFailure = { error ->
-                setStripeIntent(null)
-                onFatal(error)
-            }
+            onSuccess = ::handlePaymentSheetStateLoaded,
+            onFailure = ::handlePaymentSheetStateLoadFailure,
         )
     }
 
+    private fun handlePaymentSheetStateLoadFailure(error: Throwable) {
+        val pendingResult = pendingPaymentResult
+
+        if (pendingResult is InternalPaymentResult.Completed) {
+            // If we just received a transaction result after process death, we don't error. Instead, we dismiss
+            // PaymentSheet and return a `Completed` result to the caller.
+            val usedPaymentMethod = error.asPaymentSheetLoadingException.usedPaymentMethod
+            handlePaymentCompleted(usedPaymentMethod, finishImmediately = true)
+        } else {
+            setStripeIntent(null)
+            onFatal(error)
+        }
+
+        pendingPaymentResult = null
+    }
+
     private fun handlePaymentSheetStateLoaded(state: PaymentSheetState.Full) {
         cbcEligibility = when (state.isEligibleForCardBrandChoice) {
             true -> CardBrandChoiceEligibility.Eligible(
@@ -351,7 +365,10 @@ internal class PaymentSheetViewModel @Inject internal constructor(
 
         linkHandler.setupLink(linkState)
 
-        resetViewState()
+        val pendingFailedPaymentResult = pendingPaymentResult as? InternalPaymentResult.Failed
+        val errorMessage = pendingFailedPaymentResult?.throwable?.stripeErrorMessage(getApplication())
+
+        resetViewState(errorMessage)
         transitionToFirstScreen()
     }
 
@@ -593,74 +610,83 @@ internal class PaymentSheetViewModel @Inject internal constructor(
     }
 
     private fun onInternalPaymentResult(launcherResult: InternalPaymentResult) {
-        viewModelScope.launch {
-            runCatching {
-                requireNotNull(stripeIntent.value)
-            }.fold(
-                onSuccess = { originalIntent ->
-                    when (launcherResult) {
-                        is InternalPaymentResult.Completed -> processPayment(
-                            stripeIntent = launcherResult.intent,
-                            paymentResult = PaymentResult.Completed
-                        )
-                        is InternalPaymentResult.Failed -> processPayment(
-                            stripeIntent = originalIntent,
-                            paymentResult = PaymentResult.Failed(launcherResult.throwable)
-                        )
-                        is InternalPaymentResult.Canceled -> processPayment(
-                            stripeIntent = originalIntent,
-                            paymentResult = PaymentResult.Canceled
-                        )
-                    }
-                },
-                onFailure = ::onFatal
-            )
+        val intent = stripeIntent.value
+
+        if (intent == null) {
+            // We're recovering from process death. Wait for the pending payment result
+            // to be handled after re-loading.
+            pendingPaymentResult = launcherResult
+            return
+        }
+
+        when (launcherResult) {
+            is InternalPaymentResult.Completed -> {
+                processPayment(launcherResult.intent, PaymentResult.Completed)
+            }
+            is InternalPaymentResult.Failed -> {
+                processPayment(intent, PaymentResult.Failed(launcherResult.throwable))
+            }
+            is InternalPaymentResult.Canceled -> {
+                processPayment(intent, PaymentResult.Canceled)
+            }
+        }
+    }
+
+    private fun handlePaymentFailed(error: Throwable) {
+        eventReporter.onPaymentFailure(
+            paymentSelection = selection.value,
+            error = PaymentSheetConfirmationError.Stripe(error),
+        )
+
+        resetViewState(
+            userErrorMessage = error.stripeErrorMessage(getApplication())
+        )
+    }
+
+    private fun handlePaymentCompleted(paymentMethod: PaymentMethod?, finishImmediately: Boolean) {
+        eventReporter.onPaymentSuccess(
+            paymentSelection = selection.value,
+            deferredIntentConfirmationType = deferredIntentConfirmationType,
+        )
+
+        // Reset after sending event
+        deferredIntentConfirmationType = null
+
+        /*
+         * Sets current selection as default payment method in future payment sheet usage. New payment
+         * methods are only saved if the payment sheet is in setup mode, is in payment intent with setup
+         * for usage, or the customer has requested the payment method be saved.
+         */
+        when (val currentSelection = selection.value) {
+            is PaymentSelection.New -> paymentMethod.takeIf {
+                currentSelection.canSave(args.initializationMode)
+            }?.let { method ->
+                PaymentSelection.Saved(method)
+            }
+            else -> currentSelection
+        }?.let {
+            prefsRepository.savePaymentSelection(it)
+        }
+
+        if (finishImmediately) {
+            _paymentSheetResult.tryEmit(PaymentSheetResult.Completed)
+        } else {
+            viewState.value = PaymentSheetViewState.FinishProcessing {
+                _paymentSheetResult.tryEmit(PaymentSheetResult.Completed)
+            }
         }
     }
 
     private fun processPayment(stripeIntent: StripeIntent, paymentResult: PaymentResult) {
         when (paymentResult) {
             is PaymentResult.Completed -> {
-                eventReporter.onPaymentSuccess(
-                    paymentSelection = selection.value,
-                    deferredIntentConfirmationType = deferredIntentConfirmationType,
-                )
-
-                // Reset after sending event
-                deferredIntentConfirmationType = null
-
-                /*
-                 * Sets current selection as default payment method in future payment sheet usage. New payment
-                 * methods are only saved if the payment sheet is in setup mode, is in payment intent with setup
-                 * for usage, or the customer has requested the payment method be saved.
-                 */
-                when (val currentSelection = selection.value) {
-                    is PaymentSelection.New -> stripeIntent.paymentMethod.takeIf {
-                        currentSelection.canSave(args.initializationMode)
-                    }?.let { method ->
-                        PaymentSelection.Saved(method)
-                    }
-                    else -> currentSelection
-                }?.let {
-                    prefsRepository.savePaymentSelection(it)
-                }
-
-                viewState.value = PaymentSheetViewState.FinishProcessing {
-                    _paymentSheetResult.tryEmit(PaymentSheetResult.Completed)
-                }
+                handlePaymentCompleted(paymentMethod = stripeIntent.paymentMethod, finishImmediately = false)
             }
             is PaymentResult.Failed -> {
-                eventReporter.onPaymentFailure(
-                    paymentSelection = selection.value,
-                    error = PaymentSheetConfirmationError.Stripe(paymentResult.throwable),
-                )
-
-                resetViewState(
-                    userErrorMessage = paymentResult.throwable.stripeErrorMessage(getApplication())
-                )
+                handlePaymentFailed(paymentResult.throwable)
             }
             is PaymentResult.Canceled -> {
-                resetViewState(userErrorMessage = null)
+                resetViewState()
             }
         }
     }

paymentsheet/src/main/java/com/stripe/android/paymentsheet/model/StripeIntentValidator.kt

@@ -23,18 +23,20 @@ internal object StripeIntentValidator {
     fun requireValid(
         stripeIntent: StripeIntent
     ): StripeIntent {
+        val paymentMethod = stripeIntent.paymentMethod
+
         val exception = when {
             stripeIntent is PaymentIntent && stripeIntent.confirmationMethod != Automatic -> {
                 PaymentSheetLoadingException.InvalidConfirmationMethod(stripeIntent.confirmationMethod)
             }
             stripeIntent is PaymentIntent && stripeIntent.isInTerminalState -> {
-                PaymentSheetLoadingException.PaymentIntentInTerminalState(stripeIntent.status)
+                PaymentSheetLoadingException.PaymentIntentInTerminalState(paymentMethod, stripeIntent.status)
             }
             stripeIntent is PaymentIntent && (stripeIntent.amount == null || stripeIntent.currency == null) -> {
                 PaymentSheetLoadingException.MissingAmountOrCurrency
             }
             stripeIntent is SetupIntent && stripeIntent.isInTerminalState -> {
-                PaymentSheetLoadingException.SetupIntentInTerminalState(stripeIntent.status)
+                PaymentSheetLoadingException.SetupIntentInTerminalState(paymentMethod, stripeIntent.status)
             }
             else -> {
                 // valid

paymentsheet/src/main/java/com/stripe/android/paymentsheet/state/PaymentSheetLoadingException.kt

@@ -2,18 +2,22 @@ package com.stripe.android.paymentsheet.state
 
 import com.stripe.android.core.exception.StripeException
 import com.stripe.android.model.PaymentIntent
+import com.stripe.android.model.PaymentMethod
 import com.stripe.android.model.StripeIntent
 import com.stripe.android.paymentsheet.analytics.analyticsValue
 import com.stripe.android.paymentsheet.state.PaymentSheetLoadingException.Unknown
 
 internal sealed class PaymentSheetLoadingException : Throwable() {
 
     abstract val type: String
+    abstract val usedPaymentMethod: PaymentMethod?
 
     data class InvalidConfirmationMethod(
         private val confirmationMethod: PaymentIntent.ConfirmationMethod,
     ) : PaymentSheetLoadingException() {
 
+        override val usedPaymentMethod: PaymentMethod? = null
+
         override val type: String = "invalidConfirmationMethod"
 
         override val message: String = """
@@ -28,6 +32,8 @@ internal sealed class PaymentSheetLoadingException : Throwable() {
         private val supported: String,
     ) : PaymentSheetLoadingException() {
 
+        override val usedPaymentMethod: PaymentMethod? = null
+
         override val type: String = "noPaymentMethodTypesAvailable"
 
         override val message: String
@@ -36,6 +42,7 @@ internal sealed class PaymentSheetLoadingException : Throwable() {
     }
 
     data class PaymentIntentInTerminalState(
+        override val usedPaymentMethod: PaymentMethod?,
         private val status: StripeIntent.Status?,
     ) : PaymentSheetLoadingException() {
 
@@ -49,6 +56,7 @@ internal sealed class PaymentSheetLoadingException : Throwable() {
     }
 
     data class SetupIntentInTerminalState(
+        override val usedPaymentMethod: PaymentMethod?,
         private val status: StripeIntent.Status?,
     ) : PaymentSheetLoadingException() {
 
@@ -62,6 +70,7 @@ internal sealed class PaymentSheetLoadingException : Throwable() {
     }
 
     object MissingAmountOrCurrency : PaymentSheetLoadingException() {
+        override val usedPaymentMethod: PaymentMethod? = null
         override val type: String = "missingAmountOrCurrency"
         override val message: String = "PaymentIntent must contain amount and currency."
     }
@@ -74,6 +83,8 @@ internal sealed class PaymentSheetLoadingException : Throwable() {
             get() = StripeException.create(cause).analyticsValue
 
         override val message: String? = cause.message
+
+        override val usedPaymentMethod: PaymentMethod? = null
     }
 }
 

paymentsheet/src/test/java/com/stripe/android/paymentsheet/state/DefaultPaymentSheetLoaderTest.kt

@@ -25,8 +25,10 @@ import com.stripe.android.paymentsheet.addresselement.AddressDetails
 import com.stripe.android.paymentsheet.analytics.EventReporter
 import com.stripe.android.paymentsheet.model.PaymentSelection
 import com.stripe.android.paymentsheet.repositories.CustomerRepository
+import com.stripe.android.paymentsheet.state.PaymentSheetLoadingException.PaymentIntentInTerminalState
 import com.stripe.android.testing.FeatureFlagTestRule
 import com.stripe.android.testing.PaymentIntentFactory
+import com.stripe.android.testing.PaymentMethodFactory
 import com.stripe.android.ui.core.forms.resources.LpmRepository
 import com.stripe.android.utils.FakeCustomerRepository
 import com.stripe.android.utils.FakeElementsSessionRepository
@@ -392,18 +394,21 @@ internal class DefaultPaymentSheetLoaderTest {
 
     @Test
     fun `load() when PaymentIntent has invalid status should return null`() = runTest {
+        val paymentIntent = PaymentIntentFixtures.PI_SUCCEEDED.copy(
+            paymentMethod = PaymentMethodFactory.card(),
+        )
+        val paymentMethod = paymentIntent.paymentMethod!!
+
         val result = createPaymentSheetLoader(
-            stripeIntent = PaymentIntentFixtures.PI_REQUIRES_PAYMENT_METHOD.copy(
-                status = Succeeded,
-            ),
+            stripeIntent = paymentIntent,
         ).load(
             initializationMode = PaymentSheet.InitializationMode.PaymentIntent(
                 clientSecret = PaymentSheetFixtures.PAYMENT_INTENT_CLIENT_SECRET.value,
             ),
             PaymentSheetFixtures.CONFIG_CUSTOMER_WITH_GOOGLEPAY
         ).exceptionOrNull()
 
-        assertThat(result).isEqualTo(PaymentSheetLoadingException.PaymentIntentInTerminalState(Succeeded))
+        assertThat(result).isEqualTo(PaymentIntentInTerminalState(paymentMethod, Succeeded))
     }
 
     @Test
@@ -703,7 +708,7 @@ internal class DefaultPaymentSheetLoaderTest {
 
     @Test
     fun `Emits correct events when loading fails for deferred intent`() = runTest {
-        val error = PaymentSheetLoadingException.PaymentIntentInTerminalState(status = Canceled)
+        val error = PaymentIntentInTerminalState(usedPaymentMethod = null, status = Canceled)
         val loader = createPaymentSheetLoader(error = error)
 
         loader.load(