Prototype: in-sheet checkout session update for billing address tax

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

@@ -456,9 +456,22 @@ class Checkout private constructor(
         name: String? = null,
         phoneNumber: String? = null,
         address: Address,
+    ): Result<Unit> = updateBillingAddressInternal(
+        name = name,
+        phoneNumber = phoneNumber,
+        address = address,
+        isInSheetUpdate = false,
+    )
+
+    internal suspend fun updateBillingAddressInternal(
+        name: String?,
+        phoneNumber: String?,
+        address: Address,
+        isInSheetUpdate: Boolean,
     ): Result<Unit> {
         val built = address.build()
         return withInternalState(
+            isInSheetUpdate = isInSheetUpdate,
             additionalStateMutations = {
                 copy(billingName = name, billingPhoneNumber = phoneNumber, billingAddress = built)
             },
@@ -501,16 +514,24 @@ class Checkout private constructor(
     }
 
     private suspend fun withInternalState(
+        isInSheetUpdate: Boolean = false,
         additionalStateMutations: InternalState.() -> InternalState = { this },
         block: suspend InternalState.(sessionId: String) -> Result<CheckoutSessionResponse>,
     ): Result<Unit> {
-        if (internalState.integrationLaunched) {
+        if (!isInSheetUpdate && internalState.integrationLaunched) {
             return Result.failure(
                 IllegalStateException(
                     "Cannot mutate checkout session while a payment flow is presented."
                 )
             )
         }
+        if (internalState.checkoutSessionResponse.status != CheckoutSessionResponse.Status.OPEN) {
+            return Result.failure(
+                IllegalStateException(
+                    "Cannot mutate checkout session that is not open."
+                )
+            )
+        }
         // Run network requests with a mutex to ensure events are processed in order.
         return mutex.withLock {
             _isLoading.value = true

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

@@ -0,0 +1,69 @@
+package com.stripe.android.checkout
+
+import com.stripe.android.lpmfoundations.paymentmethod.IntegrationMetadata
+import com.stripe.android.lpmfoundations.paymentmethod.PaymentMethodMetadata
+import com.stripe.android.paymentelement.CheckoutSessionPreview
+import com.stripe.android.paymentsheet.model.PaymentSelection
+import com.stripe.android.paymentsheet.model.billingDetails
+import com.stripe.android.paymentsheet.repositories.CheckoutSessionResponse
+
+@OptIn(CheckoutSessionPreview::class)
+internal class InSheetCheckoutSessionUpdater(
+    private val checkout: Checkout?,
+) {
+    companion object {
+        fun create(paymentMethodMetadata: PaymentMethodMetadata): InSheetCheckoutSessionUpdater {
+            val metadata = paymentMethodMetadata.integrationMetadata
+                as? IntegrationMetadata.CheckoutSession
+            val checkout = metadata?.let { CheckoutInstances[it.instancesKey].firstOrNull() }
+            return InSheetCheckoutSessionUpdater(checkout)
+        }
+    }
+
+    fun requiresUpdate(): Boolean {
+        val checkout = checkout ?: return false
+        return checkout.internalState.checkoutSessionResponse.taxStatus ==
+            CheckoutSessionResponse.TaxStatus.REQUIRES_BILLING_ADDRESS
+    }
+
+    suspend fun performUpdate(
+        paymentSelection: PaymentSelection,
+    ): Result<Unit> {
+        val checkout = checkout
+            ?: return Result.failure(IllegalStateException("No Checkout instance"))
+        return updateWithBillingAddress(checkout, paymentSelection)
+    }
+
+    private suspend fun updateWithBillingAddress(
+        checkout: Checkout,
+        paymentSelection: PaymentSelection,
+    ): Result<Unit> {
+        val billingDetails = paymentSelection.billingDetails
+            ?: return Result.failure(IllegalStateException("Payment selection has no billing details"))
+
+        val billingAddress = billingDetails.address
+            ?: return Result.failure(IllegalStateException("Billing details have no address"))
+
+        if (billingAddress.country.isNullOrEmpty()) {
+            return Result.failure(IllegalStateException("Billing address has no country"))
+        }
+
+        return checkout.updateBillingAddressInternal(
+            name = billingDetails.name,
+            phoneNumber = billingDetails.phone,
+            address = billingAddress.toCheckoutAddress(),
+            isInSheetUpdate = true,
+        )
+    }
+}
+
+@OptIn(CheckoutSessionPreview::class)
+private fun com.stripe.android.model.Address.toCheckoutAddress(): Address {
+    return Address()
+        .country(requireNotNull(country))
+        .line1(line1)
+        .line2(line2)
+        .city(city)
+        .state(state)
+        .postalCode(postalCode)
+}

paymentsheet/src/main/java/com/stripe/android/paymentelement/confirmation/intent/CheckoutSessionConfirmationInterceptor.kt

@@ -102,7 +102,7 @@ internal class CheckoutSessionConfirmationInterceptor @AssistedInject constructo
             paymentMethodId = paymentMethod.id,
             clientAttributionMetadata = clientAttributionMetadata,
             returnUrl = returnUrl,
-            expectedAmount = intent.amount,
+            expectedAmount = currentExpectedAmount() ?: intent.amount,
             savePaymentMethod = savePaymentMethod,
         )
         else -> ConfirmCheckoutSessionParams(
@@ -112,6 +112,12 @@ internal class CheckoutSessionConfirmationInterceptor @AssistedInject constructo
         )
     }
 
+    private fun currentExpectedAmount(): Long? {
+        val checkout = CheckoutInstances[integrationMetadata.instancesKey].firstOrNull()
+            ?: return null
+        return checkout.internalState.checkoutSessionResponse.amount
+    }
+
     private suspend fun confirmCheckoutSession(
         params: ConfirmCheckoutSessionParams,
     ): ConfirmationDefinition.Action<Args> {

paymentsheet/src/main/java/com/stripe/android/paymentelement/embedded/sheet/SheetActivityConfirmationHelper.kt

@@ -1,5 +1,7 @@
 package com.stripe.android.paymentelement.embedded.sheet
 
+import com.stripe.android.checkout.InSheetCheckoutSessionUpdater
+import com.stripe.android.common.exception.stripeErrorMessage
 import com.stripe.android.common.model.asCommonConfiguration
 import com.stripe.android.core.injection.ViewModelScope
 import com.stripe.android.lpmfoundations.paymentmethod.PaymentMethodMetadata
@@ -32,6 +34,9 @@ internal class DefaultSheetActivityConfirmationHelper @Inject constructor(
     @ViewModelScope private val coroutineScope: CoroutineScope,
 ) : SheetActivityConfirmationHelper {
 
+    private val inSheetCheckoutSessionUpdater =
+        InSheetCheckoutSessionUpdater.create(paymentMethodMetadata)
+
     override fun confirm() {
         if (onClickDelegate.onClickOverride != null) {
             onClickDelegate.onClickOverride?.invoke()
@@ -42,13 +47,13 @@ internal class DefaultSheetActivityConfirmationHelper @Inject constructor(
 
             when (configuration.formSheetAction) {
                 EmbeddedPaymentElement.FormSheetAction.Continue -> {
-                    stateHelper.setResult(
-                        FormResult.Complete(
-                            selection = selectionHolder.selection.value,
-                            hasBeenConfirmed = false,
-                            customerState = customerStateHolder.customer.value
-                        )
-                    )
+                    if (inSheetCheckoutSessionUpdater.requiresUpdate()) {
+                        coroutineScope.launch {
+                            performCheckoutSessionUpdateThenContinue()
+                        }
+                    } else {
+                        emitContinueResult()
+                    }
                 }
                 EmbeddedPaymentElement.FormSheetAction.Confirm -> {
                     confirmationArgs()?.let { args ->
@@ -61,6 +66,34 @@ internal class DefaultSheetActivityConfirmationHelper @Inject constructor(
         }
     }
 
+    private suspend fun performCheckoutSessionUpdateThenContinue() {
+        val selection = selectionHolder.selection.value ?: run {
+            emitContinueResult()
+            return
+        }
+        stateHelper.setProcessing(true)
+        inSheetCheckoutSessionUpdater.performUpdate(selection).fold(
+            onSuccess = {
+                stateHelper.setProcessing(false)
+                emitContinueResult()
+            },
+            onFailure = { error ->
+                stateHelper.setProcessing(false)
+                stateHelper.updateError(error.stripeErrorMessage())
+            }
+        )
+    }
+
+    private fun emitContinueResult() {
+        stateHelper.setResult(
+            FormResult.Complete(
+                selection = selectionHolder.selection.value,
+                hasBeenConfirmed = false,
+                customerState = customerStateHolder.customer.value
+            )
+        )
+    }
+
     private fun confirmationArgs(): ConfirmationHandler.Args? {
         val confirmationOption = selectionHolder.selection.value?.toConfirmationOption(
             configuration = configuration.asCommonConfiguration(),

paymentsheet/src/main/java/com/stripe/android/paymentelement/embedded/sheet/SheetActivityStateHolder.kt

@@ -40,6 +40,7 @@ internal interface SheetActivityStateHolder {
     fun updateMandate(mandateText: ResolvableString?)
     fun updatePrimaryButton(callback: (PrimaryButton.UIState?) -> PrimaryButton.UIState?)
     fun updateError(error: ResolvableString?)
+    fun setProcessing(isProcessing: Boolean)
 
     fun setResult(result: FormResult)
 
@@ -167,6 +168,15 @@ internal class DefaultSheetActivityStateHolder @Inject constructor(
         }
     }
 
+    override fun setProcessing(isProcessing: Boolean) {
+        _state.update {
+            it.copy(
+                isProcessing = isProcessing,
+                isEnabled = if (isProcessing) false else selectionHolder.selection.value != null,
+            )
+        }
+    }
+
     override fun updateSavedPaymentSelectionToConfirm(selection: PaymentSelection.Saved?) {
         _state.update {
             it.copy(savedPaymentSelectionToConfirm = selection)

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

@@ -11,6 +11,7 @@ import androidx.lifecycle.viewModelScope
 import androidx.lifecycle.viewmodel.CreationExtras
 import com.stripe.android.analytics.SessionSavedStateHandler
 import com.stripe.android.cards.CardAccountRangeRepository
+import com.stripe.android.checkout.InSheetCheckoutSessionUpdater
 import com.stripe.android.common.exception.stripeErrorMessage
 import com.stripe.android.common.taptoadd.TapToAddHelper
 import com.stripe.android.common.taptoadd.TapToAddMode
@@ -51,6 +52,7 @@ import com.stripe.android.paymentsheet.ui.DefaultAddPaymentMethodInteractor
 import com.stripe.android.paymentsheet.ui.DefaultSelectSavedPaymentMethodsInteractor
 import com.stripe.android.paymentsheet.verticalmode.VerticalModeInitialScreenFactory
 import com.stripe.android.paymentsheet.viewmodels.BaseSheetViewModel
+import com.stripe.android.paymentsheet.ui.PrimaryButton
 import com.stripe.android.paymentsheet.viewmodels.PrimaryButtonUiStateMapper
 import com.stripe.android.uicore.utils.combineAsStateFlow
 import com.stripe.android.uicore.utils.mapAsStateFlow
@@ -97,6 +99,9 @@ internal class PaymentOptionsViewModel @Inject constructor(
     customViewModelScope = customViewModelScope,
 ) {
 
+    private val inSheetCheckoutSessionUpdater =
+        InSheetCheckoutSessionUpdater.create(args.state.paymentMethodMetadata)
+
     private val primaryButtonUiStateMapper = PrimaryButtonUiStateMapper(
         config = config,
         currentScreenFlow = navigationHandler.currentScreen,
@@ -386,18 +391,43 @@ internal class PaymentOptionsViewModel @Inject constructor(
                     linkAccountInfo = linkAccountHolder.linkAccountInfo.value,
                     linkExpressMode = LinkExpressMode.ENABLED,
                 )
+            } else if (inSheetCheckoutSessionUpdater.requiresUpdate()) {
+                viewModelScope.launch {
+                    performCheckoutSessionUpdateThenContinue(paymentSelection)
+                }
             } else {
-                _paymentOptionsActivityResult.tryEmit(
-                    PaymentOptionsActivityResult.Succeeded(
-                        linkAccountInfo = linkAccountHolder.linkAccountInfo.value,
-                        paymentSelection = paymentSelection.withLinkDetails(),
-                        paymentMethods = customerStateHolder.paymentMethods.value
-                    )
-                )
+                emitSucceededResult(paymentSelection)
             }
         }
     }
 
+    private suspend fun performCheckoutSessionUpdateThenContinue(paymentSelection: PaymentSelection) {
+        updatePrimaryButtonState(PrimaryButton.State.StartProcessing)
+        savedStateHandle[SAVE_PROCESSING] = true
+        inSheetCheckoutSessionUpdater.performUpdate(paymentSelection).fold(
+            onSuccess = {
+                savedStateHandle[SAVE_PROCESSING] = false
+                updatePrimaryButtonState(PrimaryButton.State.Ready)
+                emitSucceededResult(paymentSelection)
+            },
+            onFailure = { error ->
+                savedStateHandle[SAVE_PROCESSING] = false
+                updatePrimaryButtonState(PrimaryButton.State.Ready)
+                onError(error.stripeErrorMessage())
+            }
+        )
+    }
+
+    private fun emitSucceededResult(paymentSelection: PaymentSelection) {
+        _paymentOptionsActivityResult.tryEmit(
+            PaymentOptionsActivityResult.Succeeded(
+                linkAccountInfo = linkAccountHolder.linkAccountInfo.value,
+                paymentSelection = paymentSelection.withLinkDetails(),
+                paymentMethods = customerStateHolder.paymentMethods.value
+            )
+        )
+    }
+
     private fun onDisabledClick() {
         viewModelScope.launch {
             validationRequested.emit(Unit)

paymentsheet/src/main/java/com/stripe/android/paymentsheet/repositories/CheckoutSessionResponseJsonParser.kt

@@ -40,7 +40,10 @@ internal object CheckoutSessionResponseJsonParser : ModelJsonParser<CheckoutSess
         val mode = parseMode(json.optString(FIELD_MODE))
         val status = parseStatus(json.optString(FIELD_STATUS))
         val liveMode = json.optBoolean(FIELD_LIVE_MODE, false)
-        val taxStatus = parseTaxStatus(json.optJSONObject(FIELD_TAX))
+        val taxStatus = parseTaxStatusFromMeta(
+            taxMeta = json.optJSONObject(FIELD_TAX_META),
+            taxContext = json.optJSONObject(FIELD_TAX_CONTEXT),
+        )
         val amount = extractDueAmount(json) ?: return null
         val currency = json.optString(FIELD_CURRENCY).takeIf { it.isNotEmpty() } ?: return null
         val customerEmail = StripeJsonUtils.optString(json, FIELD_CUSTOMER_EMAIL)
@@ -115,14 +118,25 @@ internal object CheckoutSessionResponseJsonParser : ModelJsonParser<CheckoutSess
         }
     }
 
-    private fun parseTaxStatus(taxJson: JSONObject?): CheckoutSessionResponse.TaxStatus {
-        val statusString = taxJson?.optString(FIELD_STATUS) ?: return CheckoutSessionResponse.TaxStatus.UNKNOWN
-        return when (statusString) {
-            "ready" -> CheckoutSessionResponse.TaxStatus.READY
-            "requires_shipping_address" -> CheckoutSessionResponse.TaxStatus.REQUIRES_SHIPPING_ADDRESS
-            "requires_billing_address" -> CheckoutSessionResponse.TaxStatus.REQUIRES_BILLING_ADDRESS
-            else -> CheckoutSessionResponse.TaxStatus.UNKNOWN
+    private fun parseTaxStatusFromMeta(
+        taxMeta: JSONObject?,
+        taxContext: JSONObject?,
+    ): CheckoutSessionResponse.TaxStatus {
+        if (taxMeta == null) return CheckoutSessionResponse.TaxStatus.UNKNOWN
+        val metaStatus = taxMeta.optString(FIELD_STATUS)
+        if (metaStatus == "requires_location_inputs") {
+            val addressSource = taxContext?.optString(FIELD_AUTOMATIC_TAX_ADDRESS_SOURCE) ?: ""
+            // Default to billing (aligned with iOS and web). Only shipping if explicitly stated.
+            return if (addressSource == "session.shipping") {
+                CheckoutSessionResponse.TaxStatus.REQUIRES_SHIPPING_ADDRESS
+            } else {
+                CheckoutSessionResponse.TaxStatus.REQUIRES_BILLING_ADDRESS
+            }
+        }
+        if (metaStatus == "complete") {
+            return CheckoutSessionResponse.TaxStatus.READY
         }
+        return CheckoutSessionResponse.TaxStatus.UNKNOWN
     }
 
     private fun parseElementsSessionParams(
@@ -490,7 +504,9 @@ internal object CheckoutSessionResponseJsonParser : ModelJsonParser<CheckoutSess
     private const val FIELD_MODE = "mode"
     private const val FIELD_STATUS = "status"
     private const val FIELD_LIVE_MODE = "livemode"
-    private const val FIELD_TAX = "tax"
+    private const val FIELD_TAX_META = "tax_meta"
+    private const val FIELD_TAX_CONTEXT = "tax_context"
+    private const val FIELD_AUTOMATIC_TAX_ADDRESS_SOURCE = "automatic_tax_address_source"
     private const val FIELD_CURRENCY = "currency"
     private const val FIELD_CUSTOMER_EMAIL = "customer_email"
     private const val FIELD_ELEMENTS_SESSION = "elements_session"

paymentsheet/src/test/java/com/stripe/android/paymentelement/embedded/sheet/FakeSheetActivityStateHolder.kt

@@ -43,6 +43,10 @@ internal class FakeSheetActivityStateHolder : SheetActivityStateHolder {
         error("This should never be called!")
     }
 
+    override fun setProcessing(isProcessing: Boolean) {
+        // no-op for tests
+    }
+
     override fun setResult(result: FormResult) {
         resultTurbine.add(result)
     }