feat: plan limits

packages/ee/server-only/limits/client.ts

@@ -0,0 +1,31 @@
+import { APP_BASE_URL } from '@documenso/lib/constants/app';
+
+import { FREE_PLAN_LIMITS } from './constants';
+import { TLimitsResponseSchema, ZLimitsResponseSchema } from './schema';
+
+export type GetLimitsOptions = {
+  headers?: Record<string, string>;
+};
+
+export const getLimits = async ({ headers }: GetLimitsOptions = {}) => {
+  const requestHeaders = headers ?? {};
+
+  const url = new URL(`${APP_BASE_URL}/api/limits`);
+
+  return fetch(url, {
+    headers: {
+      ...requestHeaders,
+    },
+    next: {
+      revalidate: 60,
+    },
+  })
+    .then(async (res) => res.json())
+    .then((res) => ZLimitsResponseSchema.parse(res))
+    .catch(() => {
+      return {
+        quota: FREE_PLAN_LIMITS,
+        remaining: FREE_PLAN_LIMITS,
+      } satisfies TLimitsResponseSchema;
+    });
+};

packages/ee/server-only/limits/constants.ts

@@ -0,0 +1,11 @@
+import { TLimitsSchema } from './schema';
+
+export const FREE_PLAN_LIMITS: TLimitsSchema = {
+  documents: 5,
+  recipients: 10,
+};
+
+export const SELFHOSTED_PLAN_LIMITS: TLimitsSchema = {
+  documents: Infinity,
+  recipients: Infinity,
+};

packages/ee/server-only/limits/errors.ts

@@ -0,0 +1,6 @@
+export const ERROR_CODES: Record<string, string> = {
+  UNAUTHORIZED: 'You must be logged in to access this resource',
+  USER_FETCH_FAILED: 'An error occurred while fetching your user account',
+  SUBSCRIPTION_FETCH_FAILED: 'An error occurred while fetching your subscription',
+  UNKNOWN: 'An unknown error occurred',
+};

packages/ee/server-only/limits/handler.ts

@@ -0,0 +1,54 @@
+import { NextApiRequest, NextApiResponse } from 'next';
+
+import { getToken } from 'next-auth/jwt';
+import { match } from 'ts-pattern';
+
+import { withStaleWhileRevalidate } from '@documenso/lib/server-only/http/with-swr';
+import { getFlag } from '@documenso/lib/universal/get-feature-flag';
+
+import { SELFHOSTED_PLAN_LIMITS } from './constants';
+import { ERROR_CODES } from './errors';
+import { TLimitsErrorResponseSchema, TLimitsResponseSchema } from './schema';
+import { getServerLimits } from './server';
+
+export const limitsHandler = async (
+  req: NextApiRequest,
+  res: NextApiResponse<TLimitsResponseSchema | TLimitsErrorResponseSchema>,
+) => {
+  try {
+    const token = await getToken({ req });
+
+    const isBillingEnabled = await getFlag('app_billing');
+
+    if (!isBillingEnabled) {
+      return withStaleWhileRevalidate<typeof res>(res).status(200).json({
+        quota: SELFHOSTED_PLAN_LIMITS,
+        remaining: SELFHOSTED_PLAN_LIMITS,
+      });
+    }
+
+    if (!token?.email) {
+      throw new Error(ERROR_CODES.UNAUTHORIZED);
+    }
+
+    const limits = await getServerLimits({ email: token.email });
+
+    return withStaleWhileRevalidate<typeof res>(res).status(200).json(limits);
+  } catch (err) {
+    console.error('error', err);
+
+    if (err instanceof Error) {
+      const status = match(err.message)
+        .with(ERROR_CODES.UNAUTHORIZED, () => 401)
+        .otherwise(() => 500);
+
+      return res.status(status).json({
+        error: ERROR_CODES[err.message] ?? ERROR_CODES.UNKNOWN,
+      });
+    }
+
+    res.status(500).json({
+      error: ERROR_CODES.UNKNOWN,
+    });
+  }
+};

packages/ee/server-only/limits/provider/client.tsx

@@ -0,0 +1,53 @@
+'use client';
+
+import { createContext, useContext, useEffect, useState } from 'react';
+
+import { getLimits } from '../client';
+import { FREE_PLAN_LIMITS } from '../constants';
+import { TLimitsResponseSchema } from '../schema';
+
+export type LimitsContextValue = TLimitsResponseSchema;
+
+const LimitsContext = createContext<LimitsContextValue | null>(null);
+
+export const useLimits = () => {
+  const limits = useContext(LimitsContext);
+
+  if (!limits) {
+    throw new Error('useLimits must be used within a LimitsProvider');
+  }
+
+  return limits;
+};
+
+export type LimitsProviderProps = {
+  initialValue?: LimitsContextValue;
+  children?: React.ReactNode;
+};
+
+export const LimitsProvider = ({ initialValue, children }: LimitsProviderProps) => {
+  const defaultValue: TLimitsResponseSchema = {
+    quota: FREE_PLAN_LIMITS,
+    remaining: FREE_PLAN_LIMITS,
+  };
+
+  const [limits, setLimits] = useState(() => initialValue ?? defaultValue);
+
+  useEffect(() => {
+    void getLimits().then((limits) => setLimits(limits));
+  }, []);
+
+  useEffect(() => {
+    const onFocus = () => {
+      void getLimits().then((limits) => setLimits(limits));
+    };
+
+    window.addEventListener('focus', onFocus);
+
+    return () => {
+      window.removeEventListener('focus', onFocus);
+    };
+  }, []);
+
+  return <LimitsContext.Provider value={limits}>{children}</LimitsContext.Provider>;
+};

packages/ee/server-only/limits/provider/server.tsx

@@ -0,0 +1,18 @@
+'use server';
+
+import { headers } from 'next/headers';
+
+import { getLimits } from '../client';
+import { LimitsProvider as ClientLimitsProvider } from './client';
+
+export type LimitsProviderProps = {
+  children?: React.ReactNode;
+};
+
+export const LimitsProvider = async ({ children }: LimitsProviderProps) => {
+  const requestHeaders = Object.fromEntries(headers().entries());
+
+  const limits = await getLimits({ headers: requestHeaders });
+
+  return <ClientLimitsProvider initialValue={limits}>{children}</ClientLimitsProvider>;
+};

packages/ee/server-only/limits/schema.ts

@@ -0,0 +1,28 @@
+import { z } from 'zod';
+
+// Not proud of the below but it's a way to deal with Infinity when returning JSON.
+export const ZLimitsSchema = z.object({
+  documents: z
+    .preprocess((v) => (v === null ? Infinity : Number(v)), z.number())
+    .optional()
+    .default(0),
+  recipients: z
+    .preprocess((v) => (v === null ? Infinity : Number(v)), z.number())
+    .optional()
+    .default(0),
+});
+
+export type TLimitsSchema = z.infer<typeof ZLimitsSchema>;
+
+export const ZLimitsResponseSchema = z.object({
+  quota: ZLimitsSchema,
+  remaining: ZLimitsSchema,
+});
+
+export type TLimitsResponseSchema = z.infer<typeof ZLimitsResponseSchema>;
+
+export const ZLimitsErrorResponseSchema = z.object({
+  error: z.string(),
+});
+
+export type TLimitsErrorResponseSchema = z.infer<typeof ZLimitsErrorResponseSchema>;

packages/ee/server-only/limits/server.ts

@@ -0,0 +1,78 @@
+import { DateTime } from 'luxon';
+
+import { stripe } from '@documenso/lib/server-only/stripe';
+import { getFlag } from '@documenso/lib/universal/get-feature-flag';
+import { prisma } from '@documenso/prisma';
+
+import { FREE_PLAN_LIMITS, SELFHOSTED_PLAN_LIMITS } from './constants';
+import { ERROR_CODES } from './errors';
+import { ZLimitsSchema } from './schema';
+
+export type GetServerLimitsOptions = {
+  email?: string;
+};
+
+export const getServerLimits = async ({ email }: GetServerLimitsOptions) => {
+  const isBillingEnabled = await getFlag('app_billing');
+
+  if (!isBillingEnabled) {
+    return {
+      quota: SELFHOSTED_PLAN_LIMITS,
+      remaining: SELFHOSTED_PLAN_LIMITS,
+    };
+  }
+
+  if (!email) {
+    throw new Error(ERROR_CODES.UNAUTHORIZED);
+  }
+
+  const user = await prisma.user.findFirst({
+    where: {
+      email,
+    },
+    include: {
+      Subscription: true,
+    },
+  });
+
+  if (!user) {
+    throw new Error(ERROR_CODES.USER_FETCH_FAILED);
+  }
+
+  let quota = structuredClone(FREE_PLAN_LIMITS);
+  let remaining = structuredClone(FREE_PLAN_LIMITS);
+
+  if (user.Subscription?.priceId) {
+    const { product } = await stripe.prices
+      .retrieve(user.Subscription.priceId, {
+        expand: ['product'],
+      })
+      .catch((err) => {
+        console.error(err);
+        throw err;
+      });
+
+    if (typeof product === 'string') {
+      throw new Error(ERROR_CODES.SUBSCRIPTION_FETCH_FAILED);
+    }
+
+    quota = ZLimitsSchema.parse('metadata' in product ? product.metadata : {});
+    remaining = structuredClone(quota);
+  }
+
+  const documents = await prisma.document.count({
+    where: {
+      userId: user.id,
+      createdAt: {
+        gte: DateTime.utc().startOf('month').toJSDate(),
+      },
+    },
+  });
+
+  remaining.documents = Math.max(remaining.documents - documents, 0);
+
+  return {
+    quota,
+    remaining,
+  };
+};