Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 47 additions & 4 deletions src/composables/billing/useBillingContext.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,9 @@ const {
mockUpdateActiveWorkspace,
mockSetWorkspaceBillingRail,
mockLegacyStatus,
mockBillingStatus
mockBillingStatus,
mockActiveWorkspaceId,
mockWorkspaceTransitionGeneration
} = vi.hoisted(() => ({
mockIsPersonal: { value: true },
mockBillingRail: { value: undefined as BillingRail | undefined },
Expand All @@ -61,7 +63,9 @@ const {
subscription_tier: 'PRO',
subscription_duration: 'MONTHLY'
} as Partial<BillingStatusResponse>
}
},
mockActiveWorkspaceId: { value: 'personal-123' },
mockWorkspaceTransitionGeneration: { value: 0 }
}))

vi.mock('@vueuse/core', async (importOriginal) => {
Expand Down Expand Up @@ -90,8 +94,14 @@ vi.mock('@/platform/workspace/stores/teamWorkspaceStore', async () => {
},
get activeWorkspace() {
return mockIsPersonal.value
? { id: 'personal-123', type: 'personal' }
: { id: 'team-456', type: 'team' }
? { id: mockActiveWorkspaceId.value, type: 'personal' }
: { id: mockActiveWorkspaceId.value, type: 'team' }
},
get activeWorkspaceId() {
return mockActiveWorkspaceId.value
},
get workspaceTransitionGeneration() {
return mockWorkspaceTransitionGeneration.value
},
get activeWorkspaceBillingRail() {
return mockBillingRail.value
Expand Down Expand Up @@ -180,6 +190,8 @@ describe('useBillingContext', () => {
remoteConfig.value = {}
remoteConfigState.value = 'unloaded'
mockIsPersonal.value = true
mockActiveWorkspaceId.value = 'personal-123'
mockWorkspaceTransitionGeneration.value = 0
mockBillingRail.value = undefined
mockSetWorkspaceBillingRail.mockImplementation(
(_workspaceId: string, billingRail: BillingRail) => {
Expand Down Expand Up @@ -411,6 +423,37 @@ describe('useBillingContext', () => {
expect(mockLegacyFetchBalance).not.toHaveBeenCalled()
})

it('does not reconcile after switching away and back', async () => {
remoteConfig.value = { legacy_billing_migration_enabled: true }
remoteConfigState.value = 'authenticated'
mockBillingRail.value = 'legacy_stripe'

const context = useBillingContext()
await nextTick()
vi.clearAllMocks()

let finishStatusRefresh!: (status: BillingStatusResponse) => void
vi.mocked(workspaceApi.getBillingStatus).mockReturnValueOnce(
new Promise((resolve) => {
finishStatusRefresh = resolve
})
)

const reconciliation = context.reconcileSubscriptionSuccess()
await vi.waitFor(() =>
expect(workspaceApi.getBillingStatus).toHaveBeenCalledOnce()
)
mockActiveWorkspaceId.value = 'personal-456'
mockWorkspaceTransitionGeneration.value++
mockActiveWorkspaceId.value = 'personal-123'
mockWorkspaceTransitionGeneration.value++
finishStatusRefresh(DEFAULT_BILLING_STATUS)
await reconciliation

expect(workspaceApi.getBillingBalance).not.toHaveBeenCalled()
expect(mockLegacyFetchBalance).not.toHaveBeenCalled()
})

it('rejects topup amounts that are not positive whole-dollar cents', async () => {
const { topup } = useBillingContext()
await expect(topup(550)).rejects.toThrow()
Expand Down
18 changes: 17 additions & 1 deletion src/composables/billing/useBillingContext.ts
Original file line number Diff line number Diff line change
Expand Up @@ -285,11 +285,27 @@ function useBillingContextInternal(): BillingContext {
}

async function reconcileSubscriptionSuccess(): Promise<void> {
const workspaceId = store.activeWorkspaceId
const workspaceTransitionGeneration = store.workspaceTransitionGeneration
const checkout = checkoutContext.value
await checkout.fetchStatus()
if (
workspaceId !== store.activeWorkspaceId ||
workspaceTransitionGeneration !== store.workspaceTransitionGeneration
) {
return
}

const account = activeContext.value
if (account !== checkout) await account.fetchStatus()
if (account !== checkout) {
await account.fetchStatus()
if (
workspaceId !== store.activeWorkspaceId ||
workspaceTransitionGeneration !== store.workspaceTransitionGeneration
) {
return
}
}
await account.fetchBalance()
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,10 @@ const mockPermissions = vi.hoisted(() => ({
}))
const mockShouldUseWorkspaceBilling = vi.hoisted(() => ({ value: true }))
const mockDistributionTypes = vi.hoisted(() => ({ isCloud: true }))
const mockWorkspace = vi.hoisted(() => ({
activeWorkspaceId: 'workspace-1' as string | null,
workspaceTransitionGeneration: 0
}))

vi.mock('@/platform/distribution/types', () => mockDistributionTypes)
const mockBillingOperationState = vi.hoisted(() => ({
Expand Down Expand Up @@ -59,6 +63,10 @@ vi.mock('@/platform/workspace/stores/billingOperationStore', async () => {
}
})

vi.mock('@/platform/workspace/stores/teamWorkspaceStore', () => ({
useTeamWorkspaceStore: () => mockWorkspace
}))

vi.mock('@/platform/workspace/composables/useWorkspaceUI', async () => {
const { ref } = await import('vue')
mockPermissions.ref = ref({ canTopUp: true })
Expand Down Expand Up @@ -212,6 +220,8 @@ describe('TopUpCreditsDialogContentWorkspace', () => {
mockDistributionTypes.isCloud = true
setCanTopUp(true)
mockShouldUseWorkspaceBilling.value = true
mockWorkspace.activeWorkspaceId = 'workspace-1'
mockWorkspace.workspaceTransitionGeneration = 0
setIsAddingCredits(false)
setTopupActionOperation(undefined)
mockFetchBalance.mockResolvedValue(undefined)
Expand Down Expand Up @@ -424,6 +434,35 @@ describe('TopUpCreditsDialogContentWorkspace', () => {
})
})

it('ignores a completed top-up after switching away and back', async () => {
let resolveTopup!: (response: CreateTopupResponse) => void
mockTopup.mockReturnValueOnce(
new Promise((resolve) => {
resolveTopup = resolve
})
)

renderDialog()
await clickAddCredits()
await userEvent.click(screen.getByRole('button', { name: 'Pay $50.00' }))
await waitFor(() => expect(mockTopup).toHaveBeenCalledOnce())

mockWorkspace.activeWorkspaceId = 'workspace-2'
mockWorkspace.workspaceTransitionGeneration++
mockWorkspace.activeWorkspaceId = 'workspace-1'
mockWorkspace.workspaceTransitionGeneration++
resolveTopup(topupResponse('completed'))
await nextTick()

expect(mockFetchBalance).not.toHaveBeenCalled()
expect(mockFetchStatus).not.toHaveBeenCalled()
expect(mockToastAdd).not.toHaveBeenCalledWith(
expect.objectContaining({ severity: 'success' })
)
expect(mockCloseDialog).not.toHaveBeenCalled()
expect(mockShowSettings).not.toHaveBeenCalled()
})

it('opens Credits settings after a completed local top-up', async () => {
mockDistributionTypes.isCloud = false
mockTopup.mockResolvedValue(topupResponse('completed'))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -265,6 +265,7 @@ import { categorizeBillingApiError } from '@/platform/telemetry/utils/billingFai
import { useSettingsDialog } from '@/platform/settings/composables/useSettingsDialog'
import { useWorkspaceUI } from '@/platform/workspace/composables/useWorkspaceUI'
import { useBillingOperationStore } from '@/platform/workspace/stores/billingOperationStore'
import { useTeamWorkspaceStore } from '@/platform/workspace/stores/teamWorkspaceStore'
import { useDialogStore } from '@/stores/dialogStore'
import { cn } from '@comfyorg/tailwind-utils'

Expand All @@ -281,6 +282,7 @@ const { buildDocsUrl, docsPaths } = useExternalLink()
const { fetchBalance, fetchStatus, topup } = useBillingContext()
const { shouldUseWorkspaceBilling } = useBillingRouting()
const { permissions } = useWorkspaceUI()
const workspaceStore = useTeamWorkspaceStore()

const billingOperationStore = useBillingOperationStore()
const isPolling = computed(() => billingOperationStore.isAddingCredits)
Expand Down Expand Up @@ -425,7 +427,19 @@ async function handleBuy() {
})

const amountCents = payAmount.value * 100
const workspaceId = workspaceStore.activeWorkspaceId
const workspaceTransitionGeneration =
workspaceStore.workspaceTransitionGeneration
const response = await topup(amountCents)
if (
shouldUseWorkspaceBilling.value &&
(workspaceId !== workspaceStore.activeWorkspaceId ||
workspaceTransitionGeneration !==
workspaceStore.workspaceTransitionGeneration)
) {
paymentSubmitted.value = false
return
}
if (!response) {
paymentSubmitted.value = false
telemetry?.trackBillingEvent({
Expand Down
44 changes: 44 additions & 0 deletions src/platform/workspace/composables/useSubscriptionCheckout.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -280,13 +280,20 @@ vi.mock('@/platform/workspace/stores/billingOperationStore', () => ({
vi.mock('@/platform/workspace/stores/teamWorkspaceStore', async () => {
const { ref } = await import('vue')
const activeWorkspaceId = ref('workspace-1')
const workspaceTransitionGeneration = ref(0)
mockSetActiveWorkspaceIdImpl.value = (workspaceId) => {
if (activeWorkspaceId.value !== workspaceId) {
workspaceTransitionGeneration.value++
}
activeWorkspaceId.value = workspaceId
}
return {
useTeamWorkspaceStore: () => ({
get activeWorkspaceId() {
return activeWorkspaceId.value
},
get workspaceTransitionGeneration() {
return workspaceTransitionGeneration.value
}
})
}
Expand Down Expand Up @@ -1340,6 +1347,43 @@ describe('useSubscriptionCheckout', () => {
)
})

it('does not show synchronous success after switching away and back', async () => {
const checkout = await setup()
await checkout.handleSubscribeTeamClick({
stop: {
id: 'team_700',
usd: 700,
credits: 147_700,
discountedUsd: 665
},
billingCycle: 'monthly'
})
let resolveSubscribe!: (response: {
status: 'subscribed'
billing_op_id: string
}) => void
mockSubscribe.mockReturnValueOnce(
new Promise((resolve) => {
resolveSubscribe = resolve
})
)

const subscription = checkout.handleTeamSubscribe()
await vi.waitFor(() => expect(mockSubscribe).toHaveBeenCalledOnce())
mockSetActiveWorkspaceId('workspace-2')
mockSetActiveWorkspaceId('workspace-1')
resolveSubscribe({
status: 'subscribed',
billing_op_id: 'op-team-1'
})
await subscription

expect(checkout.checkoutStep.value).not.toBe('success')
expect(mockTrackBillingEvent).not.toHaveBeenCalledWith(
expect.objectContaining({ stage: 'succeeded' })
)
})

it('forwards confirmReactivation true when the disclosure banner reports consent', async () => {
const checkout = await setup()
await checkout.handleSubscribeTeamClick({
Expand Down
12 changes: 12 additions & 0 deletions src/platform/workspace/composables/useSubscriptionCheckout.ts
Original file line number Diff line number Diff line change
Expand Up @@ -986,6 +986,10 @@ export function useSubscriptionCheckout(
billingCycle
})
}
const workspaceIdentity = {
id: workspaceStore.activeWorkspaceId,
transitionGeneration: workspaceStore.workspaceTransitionGeneration
}
const response = await subscribe(planSlug, {
teamCreditStopId: stop.id,
billingCycle,
Expand All @@ -996,6 +1000,14 @@ export function useSubscriptionCheckout(
? previewData.value.proration_at
: undefined
})
if (
workspaceIdentity.id !== workspaceStore.activeWorkspaceId ||
workspaceIdentity.transitionGeneration !==
workspaceStore.workspaceTransitionGeneration
) {
activeCheckoutAttemptStartedAt = undefined
return
}

if (response) {
trackWorkspaceCheckoutStarted({
Expand Down
Loading
Loading