fix: preserve payment recovery rollback behavior

src/components/actionbar/ComfyRunButton/CloudRunButtonWrapper.test.ts

@@ -7,9 +7,9 @@ import CloudRunButtonWrapper from './CloudRunButtonWrapper.vue'
 
 const mockCanRunWorkflows = ref(true)
 const mockBillingStatus = ref<string | null>('paid')
+const mockV1PaymentRecovery = ref(true)
+const mockCanManageSubscription = ref(true)
 const state = vi.hoisted(() => ({
-  v1PaymentRecovery: true,
-  canManageSubscription: true,
   manageSubscription: vi.fn(),
   fetchStatus: vi.fn(),
   fetchBalance: vi.fn(),
@@ -33,7 +33,7 @@ vi.mock('@/composables/useFeatureFlags', () => ({
   useFeatureFlags: () => ({
     flags: {
       get v1PaymentRecovery() {
-        return state.v1PaymentRecovery
+        return mockV1PaymentRecovery.value
       }
     }
   })
@@ -48,7 +48,7 @@ vi.mock('@/platform/workspace/composables/useWorkspaceUI', async () => {
   return {
     useWorkspaceUI: () => ({
       permissions: computed(() => ({
-        canManageSubscription: state.canManageSubscription
+        canManageSubscription: mockCanManageSubscription.value
       }))
     })
   }
@@ -90,8 +90,8 @@ describe('CloudRunButtonWrapper', () => {
   beforeEach(() => {
     mockCanRunWorkflows.value = true
     mockBillingStatus.value = 'paid'
-    state.v1PaymentRecovery = true
-    state.canManageSubscription = true
+    mockV1PaymentRecovery.value = true
+    mockCanManageSubscription.value = true
     vi.clearAllMocks()
   })
 
@@ -196,6 +196,49 @@ describe('CloudRunButtonWrapper', () => {
     expect(state.fetchBalance).toHaveBeenCalledOnce()
   })
 
+  it('finishes the balance refresh after status unlocks recovery', async () => {
+    let resolveStatus!: () => void
+    let resolveBalance!: () => void
+    let balanceCommitted = false
+    state.fetchStatus.mockImplementationOnce(
+      (signal?: AbortSignal) =>
+        new Promise<void>((resolve) => {
+          resolveStatus = () => {
+            if (!signal?.aborted) {
+              mockBillingStatus.value = 'paid'
+              mockCanRunWorkflows.value = true
+            }
+            resolve()
+          }
+        })
+    )
+    state.fetchBalance.mockImplementationOnce(
+      (signal?: AbortSignal) =>
+        new Promise<void>((resolve) => {
+          resolveBalance = () => {
+            if (!signal?.aborted) balanceCommitted = true
+            resolve()
+          }
+        })
+    )
+    mockCanRunWorkflows.value = false
+    mockBillingStatus.value = 'paused'
+    renderWrapper()
+
+    await userEvent.click(
+      screen.getByRole('button', { name: 'Update payment to run' })
+    )
+    const dialog = state.showLayoutDialog.mock.calls[0][0]
+    await dialog.props.onUpdatePayment()
+    window.dispatchEvent(new Event('focus'))
+
+    resolveStatus()
+    await nextTick()
+    resolveBalance()
+
+    await waitFor(() => expect(balanceCommitted).toBe(true))
+  })
+
   it('reuses a pending portal request across rapid clicks and reopen', async () => {
     let resolvePortal!: () => void
     state.manageSubscription.mockImplementationOnce(
@@ -251,22 +294,37 @@ describe('CloudRunButtonWrapper', () => {
     )
     const dialog = state.showLayoutDialog.mock.calls[0][0]
     const portalRequest = dialog.props.onUpdatePayment()
+    const signal = state.manageSubscription.mock.calls[0][0]
     unmount()
 
+    expect(signal.aborted).toBe(true)
     resolvePortal()
     await portalRequest
-    expect(state.closeDialog).not.toHaveBeenCalled()
+    expect(state.closeDialog).toHaveBeenCalledOnce()
     expect(state.updateDialog).toHaveBeenCalledTimes(1)
     expect(state.updateDialog).toHaveBeenCalledWith({
       key: 'subscription-paused',
       contentProps: { isUpdatingPayment: true }
     })
   })
 
+  it('blocks a captured payment action after unmount', async () => {
+    mockCanRunWorkflows.value = false
+    mockBillingStatus.value = 'paused'
+    const { unmount } = renderWrapper()
+
+    await userEvent.click(screen.getByTestId('queue-button'))
+    const dialog = state.showLayoutDialog.mock.calls[0][0]
+    unmount()
+    await dialog.props.onUpdatePayment()
+
+    expect(state.manageSubscription).not.toHaveBeenCalled()
+  })
+
   it('opens member-safe recovery copy without a payment action', async () => {
     mockCanRunWorkflows.value = false
     mockBillingStatus.value = 'paused'
-    state.canManageSubscription = false
+    mockCanManageSubscription.value = false
     renderWrapper()
 
     expect(screen.getByTestId('queue-button')).toHaveTextContent('Run')
@@ -284,13 +342,200 @@ describe('CloudRunButtonWrapper', () => {
   it('keeps generic inactive behavior when payment recovery is disabled', () => {
     mockCanRunWorkflows.value = false
     mockBillingStatus.value = 'paused'
-    state.v1PaymentRecovery = false
+    mockV1PaymentRecovery.value = false
     renderWrapper()
 
     expect(screen.getByTestId('subscribe-to-run-button')).toBeInTheDocument()
     expect(screen.queryByTestId('queue-group')).not.toBeInTheDocument()
   })
 
+  it('closes recovery and blocks a captured payment action after rollback', async () => {
+    mockCanRunWorkflows.value = false
+    mockBillingStatus.value = 'paused'
+    renderWrapper()
+
+    await userEvent.click(screen.getByTestId('queue-button'))
+    const dialogOptions = state.showLayoutDialog.mock.calls[0][0]
+
+    mockV1PaymentRecovery.value = false
+    await nextTick()
+
+    expect(state.closeDialog).toHaveBeenCalledWith({
+      key: 'subscription-paused'
+    })
+    await dialogOptions.props.onUpdatePayment()
+    expect(state.manageSubscription).not.toHaveBeenCalled()
+  })
+
+  it('does not resume recovery UI after rollback during a portal request', async () => {
+    let resolvePortal!: () => void
+    state.manageSubscription.mockImplementationOnce(
+      () =>
+        new Promise<void>((resolve) => {
+          resolvePortal = resolve
+        })
+    )
+    mockCanRunWorkflows.value = false
+    mockBillingStatus.value = 'paused'
+    renderWrapper()
+
+    await userEvent.click(screen.getByTestId('queue-button'))
+    const dialogOptions = state.showLayoutDialog.mock.calls[0][0]
+    const portalRequest = dialogOptions.props.onUpdatePayment()
+    const signal = state.manageSubscription.mock.calls[0][0]
+
+    mockV1PaymentRecovery.value = false
+    await nextTick()
+    expect(signal.aborted).toBe(true)
+    resolvePortal()
+    await portalRequest
+    window.dispatchEvent(new Event('focus'))
+
+    expect(state.closeDialog).toHaveBeenCalledOnce()
+    expect(state.updateDialog).toHaveBeenCalledOnce()
+    expect(state.updateDialog).toHaveBeenCalledWith({
+      key: 'subscription-paused',
+      contentProps: { isUpdatingPayment: true }
+    })
+    expect(state.fetchStatus).not.toHaveBeenCalled()
+    expect(state.fetchBalance).not.toHaveBeenCalled()
+  })
+
+  it('aborts portal-return refreshes on rollback', async () => {
+    state.fetchStatus.mockReturnValueOnce(new Promise(() => {}))
+    state.fetchBalance.mockReturnValueOnce(new Promise(() => {}))
+    mockCanRunWorkflows.value = false
+    mockBillingStatus.value = 'paused'
+    renderWrapper()
+
+    await userEvent.click(screen.getByTestId('queue-button'))
+    const dialog = state.showLayoutDialog.mock.calls[0][0]
+    await dialog.props.onUpdatePayment()
+    window.dispatchEvent(new Event('focus'))
+    const statusSignal = state.fetchStatus.mock.calls[0][0]
+    const balanceSignal = state.fetchBalance.mock.calls[0][0]
+
+    mockV1PaymentRecovery.value = false
+    await nextTick()
+
+    expect(statusSignal.aborted).toBe(true)
+    expect(balanceSignal.aborted).toBe(true)
+  })
+
+  it('aborts portal-return refreshes on unmount', async () => {
+    state.fetchStatus.mockReturnValueOnce(new Promise(() => {}))
+    state.fetchBalance.mockReturnValueOnce(new Promise(() => {}))
+    mockCanRunWorkflows.value = false
+    mockBillingStatus.value = 'paused'
+    const { unmount } = renderWrapper()
+
+    await userEvent.click(screen.getByTestId('queue-button'))
+    const dialog = state.showLayoutDialog.mock.calls[0][0]
+    await dialog.props.onUpdatePayment()
+    window.dispatchEvent(new Event('focus'))
+    const statusSignal = state.fetchStatus.mock.calls[0][0]
+    const balanceSignal = state.fetchBalance.mock.calls[0][0]
+
+    unmount()
+
+    expect(statusSignal.aborted).toBe(true)
+    expect(balanceSignal.aborted).toBe(true)
+  })
+
+  it.for([
+    [
+      'owner loses permission',
+      () => {
+        mockCanManageSubscription.value = false
+      }
+    ],
+    [
+      'billing recovers',
+      () => {
+        mockBillingStatus.value = 'paid'
+      }
+    ]
+  ] as const)(
+    'invalidates payment recovery when %s',
+    async ([, invalidateRecovery]) => {
+      state.manageSubscription.mockReturnValueOnce(new Promise(() => {}))
+      mockCanRunWorkflows.value = false
+      mockBillingStatus.value = 'paused'
+      renderWrapper()
+
+      await userEvent.click(screen.getByTestId('queue-button'))
+      const dialog = state.showLayoutDialog.mock.calls[0][0]
+      void dialog.props.onUpdatePayment()
+      const signal = state.manageSubscription.mock.calls[0][0]
+
+      invalidateRecovery()
+      await nextTick()
+
+      expect(signal.aborted).toBe(true)
+      expect(state.closeDialog).toHaveBeenCalledWith({
+        key: 'subscription-paused'
+      })
+    }
+  )
+
+  it('isolates a re-enabled recovery dialog from an older portal request', async () => {
+    let resolveOldPortal!: () => void
+    state.manageSubscription.mockImplementationOnce(
+      () =>
+        new Promise<void>((resolve) => {
+          resolveOldPortal = resolve
+        })
+    )
+    mockCanRunWorkflows.value = false
+    mockBillingStatus.value = 'paused'
+    renderWrapper()
+
+    await userEvent.click(screen.getByTestId('queue-button'))
+    const oldDialog = state.showLayoutDialog.mock.calls[0][0]
+    const oldRequest = oldDialog.props.onUpdatePayment()
+
+    mockV1PaymentRecovery.value = false
+    await nextTick()
+    mockV1PaymentRecovery.value = true
+    await nextTick()
+
+    await userEvent.click(screen.getByTestId('queue-button'))
+    const newDialog = state.showLayoutDialog.mock.calls.at(-1)?.[0]
+    const newRequest = newDialog.props.onUpdatePayment()
+
+    expect(newRequest).not.toBe(oldRequest)
+    expect(state.manageSubscription).toHaveBeenCalledTimes(2)
+    await newRequest
+    const closeCount = state.closeDialog.mock.calls.length
+    const updateCount = state.updateDialog.mock.calls.length
+
+    resolveOldPortal()
+    await oldRequest
+
+    expect(state.closeDialog).toHaveBeenCalledTimes(closeCount)
+    expect(state.updateDialog).toHaveBeenCalledTimes(updateCount)
+  })
+
+  it('ignores a close callback captured by an older recovery dialog', async () => {
+    mockCanRunWorkflows.value = false
+    mockBillingStatus.value = 'paused'
+    renderWrapper()
+
+    await userEvent.click(screen.getByTestId('queue-button'))
+    const oldDialog = state.showLayoutDialog.mock.calls[0][0]
+
+    mockV1PaymentRecovery.value = false
+    await nextTick()
+    mockV1PaymentRecovery.value = true
+    await nextTick()
+    await userEvent.click(screen.getByTestId('queue-button'))
+    const closeCount = state.closeDialog.mock.calls.length
+
+    oldDialog.props.onClose()
+
+    expect(state.closeDialog).toHaveBeenCalledTimes(closeCount)
+  })
+
   it('keeps generic inactive behavior for non-paused statuses', () => {
     mockCanRunWorkflows.value = false
     mockBillingStatus.value = 'inactive'