fix: restore failed redirect checkout

Amp-Thread-ID: https://ampcode.com/threads/T-01a02c80-27c8-70bc-b916-186bbe86557a
Co-authored-by: Amp <amp@ampcode.com>

Author: huntcsg

src/platform/cloud/subscription/composables/useSubscriptionDialog.test.ts

@@ -1,5 +1,7 @@
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
+import { savePendingSubscriptionCheckout } from '@/platform/workspace/utils/pendingSubscriptionCheckout'
+
 import { useSubscriptionDialog } from './useSubscriptionDialog'
 
 const mockCloseDialog = vi.fn()
@@ -19,6 +21,10 @@ const mockIsTeamPlan = vi.hoisted(() => ({ value: false }))
 const mockCurrentPlanSlug = vi.hoisted(() => ({ value: null as string | null }))
 const mockCanManageSubscription = vi.hoisted(() => ({ value: true }))
 const mockEmbeddedCheckoutEnabled = vi.hoisted(() => ({ value: false }))
+const mockActiveWorkspaceId = vi.hoisted(() => ({ value: 'workspace-1' }))
+const mockStartOperation = vi.hoisted(() => vi.fn())
+const mockFetchPlans = vi.hoisted(() => vi.fn())
+const mockTeamCreditStops = vi.hoisted(() => ({ value: null }))
 
 vi.mock('vue', async (importOriginal) => {
   const actual = await importOriginal()
@@ -74,19 +80,28 @@ vi.mock('@/platform/distribution/types', () => ({
 
 vi.mock('@/platform/workspace/stores/teamWorkspaceStore', () => ({
   useTeamWorkspaceStore: () => ({
+    get activeWorkspaceId() {
+      return mockActiveWorkspaceId.value
+    },
     get isInPersonalWorkspace() {
       return mockIsInPersonalWorkspace.value
     }
   })
 }))
 
+vi.mock('@/platform/workspace/stores/billingOperationStore', () => ({
+  useBillingOperationStore: () => ({ startOperation: mockStartOperation })
+}))
+
 vi.mock('@/composables/billing/useBillingContext', () => ({
   useBillingContext: () => ({
     isFreeTier: mockIsFreeTier,
     isLegacyTeamPlan: mockIsLegacyTeamPlan,
     isTeamPlan: mockIsTeamPlan,
     currentPlanSlug: mockCurrentPlanSlug,
-    tier: mockTier
+    tier: mockTier,
+    fetchPlans: mockFetchPlans,
+    teamCreditStops: mockTeamCreditStops
   })
 }))
 
@@ -129,6 +144,10 @@ describe('useSubscriptionDialog', () => {
     mockCurrentPlanSlug.value = null
     mockCanManageSubscription.value = true
     mockEmbeddedCheckoutEnabled.value = false
+    mockActiveWorkspaceId.value = 'workspace-1'
+    mockStartOperation.mockResolvedValue({ status: 'succeeded' })
+    mockFetchPlans.mockResolvedValue(undefined)
+    sessionStorage.clear()
   })
 
   describe('showPricingTable', () => {
@@ -568,7 +587,7 @@ describe('useSubscriptionDialog', () => {
     it('does nothing when no resume intent is stored', () => {
       const { resumePendingPricingFlow } = useSubscriptionDialog()
 
-      resumePendingPricingFlow()
+      void resumePendingPricingFlow()
 
       expect(mockShowLayoutDialog).not.toHaveBeenCalled()
     })
@@ -580,7 +599,7 @@ describe('useSubscriptionDialog', () => {
       mockCurrentPlanSlug.value = 'creator-monthly'
 
       const { resumePendingPricingFlow } = useSubscriptionDialog()
-      resumePendingPricingFlow()
+      void resumePendingPricingFlow()
 
       expect(sessionStorage.getItem('comfy:resume-team-pricing')).toBeNull()
       expect(mockShowLayoutDialog).toHaveBeenCalledWith(
@@ -596,7 +615,7 @@ describe('useSubscriptionDialog', () => {
       mockIsInPersonalWorkspace.value = true
 
       const { resumePendingPricingFlow } = useSubscriptionDialog()
-      resumePendingPricingFlow()
+      void resumePendingPricingFlow()
 
       expect(sessionStorage.getItem('comfy:resume-team-pricing')).toBeNull()
       expect(mockShowLayoutDialog).not.toHaveBeenCalled()
@@ -607,11 +626,76 @@ describe('useSubscriptionDialog', () => {
       mockIsInPersonalWorkspace.value = false
 
       const { resumePendingPricingFlow } = useSubscriptionDialog()
-      resumePendingPricingFlow()
+      void resumePendingPricingFlow()
       mockShowLayoutDialog.mockClear()
 
-      resumePendingPricingFlow()
+      void resumePendingPricingFlow()
+      expect(mockShowLayoutDialog).not.toHaveBeenCalled()
+    })
+
+    it('reconciles a failed redirect and reopens the attempted checkout', async () => {
+      mockShouldUseWorkspaceBilling.value = true
+      mockStartOperation.mockResolvedValueOnce({ status: 'failed' })
+      savePendingSubscriptionCheckout({
+        operationId: 'op-alipay',
+        workspaceId: 'workspace-1',
+        selection: {
+          planMode: 'personal',
+          tierKey: 'creator',
+          billingCycle: 'monthly'
+        },
+        attemptedAt: Date.now()
+      })
+
+      const { resumePendingPricingFlow } = useSubscriptionDialog()
+      await resumePendingPricingFlow()
+
+      expect(mockStartOperation).toHaveBeenCalledWith(
+        'op-alipay',
+        'subscription',
+        {
+          tier: 'creator',
+          cycle: 'monthly',
+          attemptStartedAt: expect.any(Number)
+        }
+      )
+      expect(mockShowLayoutDialog).toHaveBeenCalledWith(
+        expect.objectContaining({
+          props: expect.objectContaining({
+            initialPlanMode: 'personal',
+            initialCheckout: {
+              planMode: 'personal',
+              tierKey: 'creator',
+              billingCycle: 'monthly'
+            }
+          })
+        })
+      )
+      expect(
+        sessionStorage.getItem('comfy:pending-subscription-checkout')
+      ).toBeNull()
+    })
+
+    it('clears a checkout for another workspace without reconciling it', async () => {
+      savePendingSubscriptionCheckout({
+        operationId: 'op-other-workspace',
+        workspaceId: 'workspace-2',
+        selection: {
+          planMode: 'personal',
+          tierKey: 'standard',
+          billingCycle: 'yearly'
+        },
+        attemptedAt: Date.now()
+      })
+
+      const { resumePendingPricingFlow } = useSubscriptionDialog()
+      await resumePendingPricingFlow()
+
+      expect(mockStartOperation).not.toHaveBeenCalled()
       expect(mockShowLayoutDialog).not.toHaveBeenCalled()
+      expect(
+        sessionStorage.getItem('comfy:pending-subscription-checkout')
+      ).toBeNull()
     })
   })
 })

src/platform/cloud/subscription/composables/useSubscriptionDialog.ts

@@ -4,12 +4,22 @@ import { useDialogStore } from '@/stores/dialogStore'
 import { useBillingContext } from '@/composables/billing/useBillingContext'
 import { useBillingRouting } from '@/composables/billing/useBillingRouting'
 import { useFeatureFlags } from '@/composables/useFeatureFlags'
+import {
+  getStopDiscountedMonthlyUsd,
+  mapApiTeamCreditStops
+} from '@/platform/cloud/subscription/constants/teamPlanCreditStops'
 import { isCloud } from '@/platform/distribution/types'
 import { useTelemetry } from '@/platform/telemetry'
 import type { PaymentIntentSource } from '@/platform/telemetry/types'
 import type { SubscriptionCheckoutSelection } from '@/platform/workspace/composables/useSubscriptionCheckout'
 import { useWorkspaceUI } from '@/platform/workspace/composables/useWorkspaceUI'
+import { useBillingOperationStore } from '@/platform/workspace/stores/billingOperationStore'
 import { useTeamWorkspaceStore } from '@/platform/workspace/stores/teamWorkspaceStore'
+import {
+  clearPendingSubscriptionCheckout,
+  getPendingSubscriptionCheckout
+} from '@/platform/workspace/utils/pendingSubscriptionCheckout'
+import type { PendingSubscriptionCheckout } from '@/platform/workspace/utils/pendingSubscriptionCheckout'
 
 const DIALOG_KEY = 'subscription-required'
 const RESUME_PRICING_KEY = 'comfy:resume-team-pricing'
@@ -227,11 +237,71 @@ export const useSubscriptionDialog = () => {
       })
   }
 
-  /**
-   * Check for and consume a pending team pricing resume intent.
-   * Call once after workspace initialization on app boot.
-   */
-  function resumePendingPricingFlow() {
+  async function restoreCheckoutSelection(
+    pending: PendingSubscriptionCheckout
+  ): Promise<SubscriptionCheckoutSelection | null> {
+    const selection = pending.selection
+    if (selection.planMode === 'personal') return selection
+
+    const { fetchPlans, teamCreditStops } = useBillingContext()
+    await fetchPlans()
+    const stop = mapApiTeamCreditStops(teamCreditStops.value?.stops ?? []).find(
+      ({ id }) => id === selection.teamCreditStopId
+    )
+    if (!stop?.id) return null
+
+    return {
+      planMode: 'team',
+      stop: {
+        id: stop.id,
+        usd: stop.usd,
+        credits: stop.credits,
+        discountedUsd: getStopDiscountedMonthlyUsd(stop, selection.billingCycle)
+      },
+      billingCycle: selection.billingCycle
+    }
+  }
+
+  async function resumePendingCheckout(
+    pending: PendingSubscriptionCheckout
+  ): Promise<void> {
+    if (pending.workspaceId !== workspaceStore.activeWorkspaceId) {
+      clearPendingSubscriptionCheckout(pending.operationId)
+      return
+    }
+
+    const billingOperationStore = useBillingOperationStore()
+    try {
+      const operation = await billingOperationStore.startOperation(
+        pending.operationId,
+        'subscription',
+        {
+          tier:
+            pending.selection.planMode === 'personal'
+              ? pending.selection.tierKey
+              : 'team',
+          cycle: pending.selection.billingCycle,
+          attemptStartedAt: pending.attemptedAt
+        }
+      )
+      if (operation.status === 'succeeded') return
+
+      const initialCheckout = await restoreCheckoutSelection(pending)
+      if (initialCheckout) {
+        showPricingTable({
+          planMode: initialCheckout.planMode,
+          initialCheckout
+        })
+      }
+    } finally {
+      clearPendingSubscriptionCheckout(pending.operationId)
+    }
+  }
+
+  function resumePendingPricingFlow(): Promise<void> | void {
+    const pendingCheckout = getPendingSubscriptionCheckout()
+    if (pendingCheckout) return resumePendingCheckout(pendingCheckout)
+
     try {
       const pending = sessionStorage.getItem(RESUME_PRICING_KEY)
       if (!pending) return

src/platform/workspace/auth/WorkspaceAuthGate.vue

@@ -173,12 +173,10 @@ async function initialize(): Promise<void> {
       throw new Error('Unified cloud auth was cleared during workspace setup')
     }
 
-    // Resume any pending pricing flow from team workspace creation
-    // Only safe after workspace store initialized successfully — the pricing
-    // dialog reads workspace state to decide which variant to show.
+    // Pricing recovery reads the initialized workspace to scope its intent.
     const workspaceStore = useTeamWorkspaceStore()
     if (workspaceStore.initState === 'ready') {
-      subscriptionDialog.resumePendingPricingFlow()
+      void subscriptionDialog.resumePendingPricingFlow()
     }
 
     if (generation === initializationGeneration) {

src/platform/workspace/composables/useSubscriptionCheckout.test.ts

@@ -463,6 +463,7 @@ describe('useSubscriptionCheckout', () => {
       canDowngradeToPersonal: true
     }
     mockSubscription.value = null
+    sessionStorage.clear()
     emit = vi.fn()
   })
 
@@ -2962,6 +2963,51 @@ describe('useSubscriptionCheckout', () => {
       expect(checkout.checkoutStep.value).toBe('preview')
     })
 
+    it('persists the pending attempt until its operation becomes terminal', async () => {
+      const checkout = await setup()
+      checkout.selectedTierKey.value = 'creator'
+      checkout.selectedBillingCycle.value = 'monthly'
+      mockSubscribe.mockResolvedValueOnce({
+        status: 'pending_payment',
+        billing_op_id: 'op-alipay'
+      })
+      let resolveOperation!: (operation: {
+        status: 'failed'
+        workspaceId: string
+      }) => void
+      mockStartOperation.mockImplementationOnce(
+        () =>
+          new Promise((resolve) => {
+            resolveOperation = resolve
+          })
+      )
+
+      const payment = checkout.handleAddCreditCard()
+      await vi.waitFor(() => {
+        expect(
+          JSON.parse(
+            sessionStorage.getItem('comfy:pending-subscription-checkout') ??
+              'null'
+          )
+        ).toMatchObject({
+          operationId: 'op-alipay',
+          workspaceId: 'workspace-1',
+          selection: {
+            planMode: 'personal',
+            tierKey: 'creator',
+            billingCycle: 'monthly'
+          }
+        })
+      })
+
+      resolveOperation({ status: 'failed', workspaceId: 'workspace-1' })
+      await payment
+
+      expect(
+        sessionStorage.getItem('comfy:pending-subscription-checkout')
+      ).toBeNull()
+    })
+
     it('shows error toast on subscribe failure', async () => {
       const checkout = await setupWithApprovedPreview()
       checkout.selectedTierKey.value = 'standard'