feat: add idempotent request support for payment apis (#4753)

packages/backend/server/src/modules/payment/resolver.ts

@@ -189,12 +189,14 @@ export class SubscriptionResolver {
   async checkout(
     @CurrentUser() user: User,
     @Args({ name: 'recurring', type: () => SubscriptionRecurring })
-    recurring: SubscriptionRecurring
+    recurring: SubscriptionRecurring,
+    @Args('idempotencyKey') idempotencyKey: string
   ) {
     const session = await this.service.createCheckoutSession({
       user,
       recurring,
       redirectUrl: `${this.config.baseUrl}/upgrade-success`,
+      idempotencyKey,
     });
 
     if (!session.url) {
@@ -217,22 +219,33 @@ export class SubscriptionResolver {
   }
 
   @Mutation(() => UserSubscriptionType)
-  async cancelSubscription(@CurrentUser() user: User) {
-    return this.service.cancelSubscription(user.id);
+  async cancelSubscription(
+    @CurrentUser() user: User,
+    @Args('idempotencyKey') idempotencyKey: string
+  ) {
+    return this.service.cancelSubscription(idempotencyKey, user.id);
   }
 
   @Mutation(() => UserSubscriptionType)
-  async resumeSubscription(@CurrentUser() user: User) {
-    return this.service.resumeCanceledSubscription(user.id);
+  async resumeSubscription(
+    @CurrentUser() user: User,
+    @Args('idempotencyKey') idempotencyKey: string
+  ) {
+    return this.service.resumeCanceledSubscription(idempotencyKey, user.id);
   }
 
   @Mutation(() => UserSubscriptionType)
   async updateSubscriptionRecurring(
     @CurrentUser() user: User,
     @Args({ name: 'recurring', type: () => SubscriptionRecurring })
-    recurring: SubscriptionRecurring
+    recurring: SubscriptionRecurring,
+    @Args('idempotencyKey') idempotencyKey: string
   ) {
-    return this.service.updateSubscriptionRecurring(user.id, recurring);
+    return this.service.updateSubscriptionRecurring(
+      idempotencyKey,
+      user.id,
+      recurring
+    );
   }
 }
 

packages/backend/server/src/modules/payment/schedule.ts

@@ -1,9 +1,10 @@
-import { Injectable } from '@nestjs/common';
+import { Injectable, Logger } from '@nestjs/common';
 import Stripe from 'stripe';
 
 @Injectable()
 export class ScheduleManager {
   private _schedule: Stripe.SubscriptionSchedule | null = null;
+  private readonly logger = new Logger(ScheduleManager.name);
 
   constructor(private readonly stripe: Stripe) {}
 
@@ -50,7 +51,10 @@ export class ScheduleManager {
     if (typeof schedule === 'string') {
       const s = await this.stripe.subscriptionSchedules
         .retrieve(schedule)
-        .catch(() => undefined);
+        .catch(e => {
+          this.logger.error('Failed to retrieve subscription schedule', e);
+          return undefined;
+        });
 
       return ScheduleManager.create(this.stripe, s);
     } else {
@@ -58,7 +62,10 @@ export class ScheduleManager {
     }
   }
 
-  async fromSubscription(subscription: string | Stripe.Subscription) {
+  async fromSubscription(
+    idempotencyKey: string,
+    subscription: string | Stripe.Subscription
+  ) {
     if (typeof subscription === 'string') {
       subscription = await this.stripe.subscriptions.retrieve(subscription, {
         expand: ['schedule'],
@@ -68,9 +75,10 @@ export class ScheduleManager {
     if (subscription.schedule) {
       return await this.fromSchedule(subscription.schedule);
     } else {
-      const schedule = await this.stripe.subscriptionSchedules.create({
-        from_subscription: subscription.id,
-      });
+      const schedule = await this.stripe.subscriptionSchedules.create(
+        { from_subscription: subscription.id },
+        { idempotencyKey }
+      );
 
       return await this.fromSchedule(schedule);
     }
@@ -80,7 +88,7 @@ export class ScheduleManager {
    * Cancel a subscription by marking schedule's end behavior to `cancel`.
    * At the same time, the coming phase's price and coupon will be saved to metadata for later resuming to correction subscription.
    */
-  async cancel() {
+  async cancel(idempotencyKey: string) {
     if (!this._schedule) {
       throw new Error('No schedule');
     }
@@ -111,13 +119,17 @@ export class ScheduleManager {
       };
     }
 
-    await this.stripe.subscriptionSchedules.update(this._schedule.id, {
-      phases: [phases],
-      end_behavior: 'cancel',
-    });
+    await this.stripe.subscriptionSchedules.update(
+      this._schedule.id,
+      {
+        phases: [phases],
+        end_behavior: 'cancel',
+      },
+      { idempotencyKey }
+    );
   }
 
-  async resume() {
+  async resume(idempotencyKey: string) {
     if (!this._schedule) {
       throw new Error('No schedule');
     }
@@ -156,21 +168,27 @@ export class ScheduleManager {
       });
     }
 
-    await this.stripe.subscriptionSchedules.update(this._schedule.id, {
-      phases: phases,
-      end_behavior: 'release',
-    });
+    await this.stripe.subscriptionSchedules.update(
+      this._schedule.id,
+      {
+        phases: phases,
+        end_behavior: 'release',
+      },
+      { idempotencyKey }
+    );
   }
 
-  async release() {
+  async release(idempotencyKey: string) {
     if (!this._schedule) {
       throw new Error('No schedule');
     }
 
-    await this.stripe.subscriptionSchedules.release(this._schedule.id);
+    await this.stripe.subscriptionSchedules.release(this._schedule.id, {
+      idempotencyKey,
+    });
   }
 
-  async update(price: string, coupon?: string) {
+  async update(idempotencyKey: string, price: string, coupon?: string) {
     if (!this._schedule) {
       throw new Error('No schedule');
     }
@@ -184,31 +202,37 @@ export class ScheduleManager {
       this.currentPhase.items[0].price === price &&
       (!coupon || this.currentPhase.coupon === coupon)
     ) {
-      await this.stripe.subscriptionSchedules.release(this._schedule.id);
+      await this.stripe.subscriptionSchedules.release(this._schedule.id, {
+        idempotencyKey,
+      });
       this._schedule = null;
     } else {
-      await this.stripe.subscriptionSchedules.update(this._schedule.id, {
-        phases: [
-          {
-            items: [
-              {
-                price: this.currentPhase.items[0].price as string,
-              },
-            ],
-            start_date: this.currentPhase.start_date,
-            end_date: this.currentPhase.end_date,
-          },
-          {
-            items: [
-              {
-                price: price,
-                quantity: 1,
-              },
-            ],
-            coupon,
-          },
-        ],
-      });
+      await this.stripe.subscriptionSchedules.update(
+        this._schedule.id,
+        {
+          phases: [
+            {
+              items: [
+                {
+                  price: this.currentPhase.items[0].price as string,
+                },
+              ],
+              start_date: this.currentPhase.start_date,
+              end_date: this.currentPhase.end_date,
+            },
+            {
+              items: [
+                {
+                  price: price,
+                  quantity: 1,
+                },
+              ],
+              coupon,
+            },
+          ],
+        },
+        { idempotencyKey }
+      );
     }
   }
 }

packages/backend/server/src/modules/payment/service.ts

@@ -103,12 +103,14 @@ export class SubscriptionService {
     user,
     recurring,
     redirectUrl,
+    idempotencyKey,
     plan = SubscriptionPlan.Pro,
   }: {
     user: User;
     plan?: SubscriptionPlan;
     recurring: SubscriptionRecurring;
     redirectUrl: string;
+    idempotencyKey: string;
   }) {
     const currentSubscription = await this.db.userSubscription.findUnique({
       where: {
@@ -121,37 +123,43 @@ export class SubscriptionService {
     }
 
     const price = await this.getPrice(plan, recurring);
-    const customer = await this.getOrCreateCustomer(user);
+    const customer = await this.getOrCreateCustomer(idempotencyKey, user);
     const coupon = await this.getAvailableCoupon(user, CouponType.EarlyAccess);
 
-    return await this.stripe.checkout.sessions.create({
-      line_items: [
-        {
-          price,
-          quantity: 1,
+    return await this.stripe.checkout.sessions.create(
+      {
+        line_items: [
+          {
+            price,
+            quantity: 1,
+          },
+        ],
+        tax_id_collection: {
+          enabled: true,
+        },
+        ...(coupon
+          ? {
+              discounts: [{ coupon }],
+            }
+          : {
+              allow_promotion_codes: true,
+            }),
+        mode: 'subscription',
+        success_url: redirectUrl,
+        customer: customer.stripeCustomerId,
+        customer_update: {
+          address: 'auto',
+          name: 'auto',
         },
-      ],
-      tax_id_collection: {
-        enabled: true,
       },
-      ...(coupon
-        ? {
-            discounts: [{ coupon }],
-          }
-        : {
-            allow_promotion_codes: true,
-          }),
-      mode: 'subscription',
-      success_url: redirectUrl,
-      customer: customer.stripeCustomerId,
-      customer_update: {
-        address: 'auto',
-        name: 'auto',
-      },
-    });
+      { idempotencyKey }
+    );
   }
 
-  async cancelSubscription(userId: string): Promise<UserSubscription> {
+  async cancelSubscription(
+    idempotencyKey: string,
+    userId: string
+  ): Promise<UserSubscription> {
     const user = await this.db.user.findUnique({
       where: {
         id: userId,
@@ -174,7 +182,7 @@ export class SubscriptionService {
       const manager = await this.scheduleManager.fromSchedule(
         user.subscription.stripeScheduleId
       );
-      await manager.cancel();
+      await manager.cancel(idempotencyKey);
       return this.saveSubscription(
         user,
         await this.stripe.subscriptions.retrieve(
@@ -187,15 +195,17 @@ export class SubscriptionService {
       // see https://stripe.com/docs/billing/subscriptions/cancel
       const subscription = await this.stripe.subscriptions.update(
         user.subscription.stripeSubscriptionId,
-        {
-          cancel_at_period_end: true,
-        }
+        { cancel_at_period_end: true },
+        { idempotencyKey }
       );
       return await this.saveSubscription(user, subscription);
     }
   }
 
-  async resumeCanceledSubscription(userId: string): Promise<UserSubscription> {
+  async resumeCanceledSubscription(
+    idempotencyKey: string,
+    userId: string
+  ): Promise<UserSubscription> {
     const user = await this.db.user.findUnique({
       where: {
         id: userId,
@@ -221,7 +231,7 @@ export class SubscriptionService {
       const manager = await this.scheduleManager.fromSchedule(
         user.subscription.stripeScheduleId
       );
-      await manager.resume();
+      await manager.resume(idempotencyKey);
       return this.saveSubscription(
         user,
         await this.stripe.subscriptions.retrieve(
@@ -232,9 +242,8 @@ export class SubscriptionService {
     } else {
       const subscription = await this.stripe.subscriptions.update(
         user.subscription.stripeSubscriptionId,
-        {
-          cancel_at_period_end: false,
-        }
+        { cancel_at_period_end: false },
+        { idempotencyKey }
       );
 
       return await this.saveSubscription(user, subscription);
@@ -242,6 +251,7 @@ export class SubscriptionService {
   }
 
   async updateSubscriptionRecurring(
+    idempotencyKey: string,
     userId: string,
     recurring: SubscriptionRecurring
   ): Promise<UserSubscription> {
@@ -272,10 +282,12 @@ export class SubscriptionService {
     );
 
     const manager = await this.scheduleManager.fromSubscription(
+      idempotencyKey,
       user.subscription.stripeSubscriptionId
     );
 
     await manager.update(
+      idempotencyKey,
       price,
       // if user is early access user, use early access coupon
       manager.currentPhase?.coupon === CouponType.EarlyAccess ||
@@ -355,10 +367,16 @@ export class SubscriptionService {
 
     // deal with early access user
     if (stripeInvoice.discount?.coupon.id === CouponType.EarlyAccess) {
+      const idempotencyKey = stripeInvoice.id + '_earlyaccess';
       const manager = await this.scheduleManager.fromSubscription(
+        idempotencyKey,
         line.subscription as string
       );
-      await manager.update(line.price.id, CouponType.EarlyAccessRenew);
+      await manager.update(
+        idempotencyKey,
+        line.price.id,
+        CouponType.EarlyAccessRenew
+      );
     }
   }
 
@@ -530,7 +548,10 @@ export class SubscriptionService {
     }
   }
 
-  private async getOrCreateCustomer(user: User): Promise<UserStripeCustomer> {
+  private async getOrCreateCustomer(
+    idempotencyKey: string,
+    user: User
+  ): Promise<UserStripeCustomer> {
     const customer = await this.db.userStripeCustomer.findUnique({
       where: {
         userId: user.id,
@@ -550,9 +571,10 @@ export class SubscriptionService {
     if (stripeCustomersList.data.length) {
       stripeCustomer = stripeCustomersList.data[0];
     } else {
-      stripeCustomer = await this.stripe.customers.create({
-        email: user.email,
-      });
+      stripeCustomer = await this.stripe.customers.create(
+        { email: user.email },
+        { idempotencyKey }
+      );
     }
 
     return await this.db.userStripeCustomer.create({

packages/backend/server/src/schema.gql

@@ -278,13 +278,13 @@ type Mutation {
   addToNewFeaturesWaitingList(type: NewFeaturesKind!, email: String!): AddToNewFeaturesWaitingList!
 
   """Create a subscription checkout link of stripe"""
-  checkout(recurring: SubscriptionRecurring!): String!
+  checkout(recurring: SubscriptionRecurring!, idempotencyKey: String!): String!
 
   """Create a stripe customer portal to manage payment methods"""
   createCustomerPortal: String!
-  cancelSubscription: UserSubscription!
-  resumeSubscription: UserSubscription!
-  updateSubscriptionRecurring(recurring: SubscriptionRecurring!): UserSubscription!
+  cancelSubscription(idempotencyKey: String!): UserSubscription!
+  resumeSubscription(idempotencyKey: String!): UserSubscription!
+  updateSubscriptionRecurring(recurring: SubscriptionRecurring!, idempotencyKey: String!): UserSubscription!
 }
 
 """The `Upload` scalar type represents a file upload."""

packages/frontend/core/src/components/affine/setting-modal/general-setting/billing/index.tsx

@@ -21,6 +21,7 @@ import { useMutation, useQuery } from '@affine/workspace/affine/gql';
 import { ArrowRightSmallIcon } from '@blocksuite/icons';
 import { Button, IconButton } from '@toeverything/components/button';
 import { useSetAtom } from 'jotai';
+import { nanoid } from 'nanoid';
 import { Suspense, useCallback, useMemo, useState } from 'react';
 
 import { openSettingModalAtom } from '../../../../../atoms';
@@ -89,13 +90,19 @@ const SubscriptionSettings = () => {
   });
   const [openCancelModal, setOpenCancelModal] = useState(false);
 
+  // allow replay request on network error until component unmount
+  const idempotencyKey = useMemo(() => nanoid(), []);
+
   const cancel = useCallback(() => {
-    trigger(null, {
-      onSuccess: data => {
-        mutateSubscription(data.cancelSubscription);
-      },
-    });
-  }, [trigger, mutateSubscription]);
+    trigger(
+      { idempotencyKey },
+      {
+        onSuccess: data => {
+          mutateSubscription(data.cancelSubscription);
+        },
+      }
+    );
+  }, [trigger, idempotencyKey, mutateSubscription]);
 
   const { data: pricesQueryResult } = useQuery({
     query: pricesQuery,
@@ -288,13 +295,19 @@ const ResumeSubscription = ({
     mutation: resumeSubscriptionMutation,
   });
 
+  // allow replay request on network error until component unmount
+  const idempotencyKey = useMemo(() => nanoid(), []);
+
   const resume = useCallback(() => {
-    trigger(null, {
-      onSuccess: data => {
-        onSubscriptionUpdate(data.resumeSubscription);
-      },
-    });
-  }, [trigger, onSubscriptionUpdate]);
+    trigger(
+      { idempotencyKey },
+      {
+        onSuccess: data => {
+          onSubscriptionUpdate(data.resumeSubscription);
+        },
+      }
+    );
+  }, [trigger, idempotencyKey, onSubscriptionUpdate]);
 
   return (
     <Button

packages/frontend/core/src/components/affine/setting-modal/general-setting/plans/plan-card.tsx

@@ -18,10 +18,12 @@ import { Button } from '@toeverything/components/button';
 import { Tooltip } from '@toeverything/components/tooltip';
 import { useSetAtom } from 'jotai';
 import { useAtom } from 'jotai';
+import { nanoid } from 'nanoid';
 import {
   type PropsWithChildren,
   useCallback,
   useEffect,
+  useMemo,
   useRef,
   useState,
 } from 'react';
@@ -303,13 +305,19 @@ const Downgrade = ({
     mutation: cancelSubscriptionMutation,
   });
 
+  // allow replay request on network error until component unmount
+  const idempotencyKey = useMemo(() => nanoid(), []);
+
   const downgrade = useCallback(() => {
-    trigger(null, {
-      onSuccess: data => {
-        onSubscriptionUpdate(data.cancelSubscription);
-      },
-    });
-  }, [trigger, onSubscriptionUpdate]);
+    trigger(
+      { idempotencyKey },
+      {
+        onSuccess: data => {
+          onSubscriptionUpdate(data.cancelSubscription);
+        },
+      }
+    );
+  }, [trigger, idempotencyKey, onSubscriptionUpdate]);
 
   const tooltipContent = disabled
     ? t['com.affine.payment.downgraded-tooltip']()
@@ -365,6 +373,9 @@ const Upgrade = ({
 
   const newTabRef = useRef<Window | null>(null);
 
+  // allow replay request on network error until component unmount
+  const idempotencyKey = useMemo(() => nanoid(), []);
+
   const onClose = useCallback(() => {
     newTabRef.current = null;
     onSubscriptionUpdate();
@@ -381,7 +392,7 @@ const Upgrade = ({
       newTabRef.current.focus();
     } else {
       trigger(
-        { recurring },
+        { recurring, idempotencyKey },
         {
           onSuccess: data => {
             // FIXME: safari prevents from opening new tab by window api
@@ -401,7 +412,7 @@ const Upgrade = ({
         }
       );
     }
-  }, [trigger, recurring, onClose, openPaymentDisableModal]);
+  }, [openPaymentDisableModal, trigger, recurring, idempotencyKey, onClose]);
 
   useEffect(() => {
     return () => {
@@ -446,16 +457,19 @@ const ChangeRecurring = ({
     mutation: updateSubscriptionMutation,
   });
 
+  // allow replay request on network error until component unmount
+  const idempotencyKey = useMemo(() => nanoid(), []);
+
   const change = useCallback(() => {
     trigger(
-      { recurring: to },
+      { recurring: to, idempotencyKey },
       {
         onSuccess: data => {
           onSubscriptionUpdate(data.updateSubscriptionRecurring);
         },
       }
     );
-  }, [trigger, onSubscriptionUpdate, to]);
+  }, [trigger, to, idempotencyKey, onSubscriptionUpdate]);
 
   const changeCurringContent = (
     <Trans values={{ from, to, due }} className={styles.downgradeContent}>
@@ -523,13 +537,19 @@ const ResumeAction = ({
     mutation: resumeSubscriptionMutation,
   });
 
+  // allow replay request on network error until component unmount
+  const idempotencyKey = useMemo(() => nanoid(), []);
+
   const resume = useCallback(() => {
-    trigger(null, {
-      onSuccess: data => {
-        onSubscriptionUpdate(data.resumeSubscription);
-      },
-    });
-  }, [trigger, onSubscriptionUpdate]);
+    trigger(
+      { idempotencyKey },
+      {
+        onSuccess: data => {
+          onSubscriptionUpdate(data.resumeSubscription);
+        },
+      }
+    );
+  }, [trigger, idempotencyKey, onSubscriptionUpdate]);
 
   return (
     <>

packages/frontend/graphql/src/graphql/cancel-subscription.gql

@@ -1,5 +1,5 @@
-mutation cancelSubscription {
-  cancelSubscription {
+mutation cancelSubscription($idempotencyKey: String!) {
+  cancelSubscription(idempotencyKey: $idempotencyKey) {
     id
     status
     nextBillAt

packages/frontend/graphql/src/graphql/create-checkout-link.gql

@@ -1,3 +1,6 @@
-mutation checkout($recurring: SubscriptionRecurring!) {
-  checkout(recurring: $recurring)
+mutation checkout(
+  $recurring: SubscriptionRecurring!
+  $idempotencyKey: String!
+) {
+  checkout(recurring: $recurring, idempotencyKey: $idempotencyKey)
 }

packages/frontend/graphql/src/graphql/index.ts

@@ -85,8 +85,8 @@ export const cancelSubscriptionMutation = {
   definitionName: 'cancelSubscription',
   containsFile: false,
   query: `
-mutation cancelSubscription {
-  cancelSubscription {
+mutation cancelSubscription($idempotencyKey: String!) {
+  cancelSubscription(idempotencyKey: $idempotencyKey) {
     id
     status
     nextBillAt
@@ -133,8 +133,8 @@ export const checkoutMutation = {
   definitionName: 'checkout',
   containsFile: false,
   query: `
-mutation checkout($recurring: SubscriptionRecurring!) {
-  checkout(recurring: $recurring)
+mutation checkout($recurring: SubscriptionRecurring!, $idempotencyKey: String!) {
+  checkout(recurring: $recurring, idempotencyKey: $idempotencyKey)
 }`,
 };
 
@@ -434,8 +434,8 @@ export const resumeSubscriptionMutation = {
   definitionName: 'resumeSubscription',
   containsFile: false,
   query: `
-mutation resumeSubscription {
-  resumeSubscription {
+mutation resumeSubscription($idempotencyKey: String!) {
+  resumeSubscription(idempotencyKey: $idempotencyKey) {
     id
     status
     nextBillAt
@@ -593,8 +593,11 @@ export const updateSubscriptionMutation = {
   definitionName: 'updateSubscriptionRecurring',
   containsFile: false,
   query: `
-mutation updateSubscription($recurring: SubscriptionRecurring!) {
-  updateSubscriptionRecurring(recurring: $recurring) {
+mutation updateSubscription($recurring: SubscriptionRecurring!, $idempotencyKey: String!) {
+  updateSubscriptionRecurring(
+    recurring: $recurring
+    idempotencyKey: $idempotencyKey
+  ) {
     id
     plan
     recurring

packages/frontend/graphql/src/graphql/resume-subscription.gql

@@ -1,5 +1,5 @@
-mutation resumeSubscription {
-  resumeSubscription {
+mutation resumeSubscription($idempotencyKey: String!) {
+  resumeSubscription(idempotencyKey: $idempotencyKey) {
     id
     status
     nextBillAt

packages/frontend/graphql/src/graphql/update-subscription-billing.gql

@@ -1,5 +1,11 @@
-mutation updateSubscription($recurring: SubscriptionRecurring!) {
-  updateSubscriptionRecurring(recurring: $recurring) {
+mutation updateSubscription(
+  $recurring: SubscriptionRecurring!
+  $idempotencyKey: String!
+) {
+  updateSubscriptionRecurring(
+    recurring: $recurring
+    idempotencyKey: $idempotencyKey
+  ) {
     id
     plan
     recurring

packages/frontend/graphql/src/schema.ts

@@ -131,7 +131,7 @@ export type AllBlobSizesQuery = {
 };
 
 export type CancelSubscriptionMutationVariables = Exact<{
-  [key: string]: never;
+  idempotencyKey: Scalars['String']['input'];
 }>;
 
 export type CancelSubscriptionMutation = {
@@ -178,6 +178,7 @@ export type ChangePasswordMutation = {
 
 export type CheckoutMutationVariables = Exact<{
   recurring: SubscriptionRecurring;
+  idempotencyKey: Scalars['String']['input'];
 }>;
 
 export type CheckoutMutation = { __typename?: 'Mutation'; checkout: string };
@@ -416,7 +417,7 @@ export type RemoveAvatarMutation = {
 };
 
 export type ResumeSubscriptionMutationVariables = Exact<{
-  [key: string]: never;
+  idempotencyKey: Scalars['String']['input'];
 }>;
 
 export type ResumeSubscriptionMutation = {
@@ -558,6 +559,7 @@ export type SubscriptionQuery = {
 
 export type UpdateSubscriptionMutationVariables = Exact<{
   recurring: SubscriptionRecurring;
+  idempotencyKey: Scalars['String']['input'];
 }>;
 
 export type UpdateSubscriptionMutation = {