Skip to content

Commit

Permalink
Reduce number of static subscription functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@frankaloia
frankaloia committed Mar 12, 2025
1 parent 483fea7 commit 24dc499
Show file tree
Hide file tree
Showing 5 changed files with 106 additions and 122 deletions.
11 changes: 11 additions & 0 deletions front/lib/auth.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -736,6 +736,17 @@ export class Authenticator {
return this._subscription;
}

getNonNullableSubscriptionResource(): SubscriptionResource {
const subscriptionResource = this.subscriptionResource();

if (!subscriptionResource) {
throw new Error("Unexpected unauthenticated call to `getNonNullableSubscriptionResource`.");
}

return subscriptionResource;
}


plan(): PlanType | null {
return this._subscription ? this._subscription.getPlan() : null;
}
Expand Down
17 changes: 5 additions & 12 deletions front/lib/plans/stripe.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,13 @@ import type {
LightWorkspaceType,
Result,
SubscriptionType,
UserType,
WorkspaceType,
} from "@dust-tt/types";
import { Err, isDevelopment, Ok } from "@dust-tt/types";
import { Stripe } from "stripe";

import config from "@app/lib/api/config";
import type { Authenticator } from "@app/lib/auth";
import { Plan, Subscription } from "@app/lib/models/plan";
import { isOldFreePlan } from "@app/lib/plans/plan_codes";
import { countActiveSeatsInWorkspace } from "@app/lib/plans/usage/seats";
Expand Down Expand Up @@ -77,25 +77,18 @@ async function getDefautPriceFromMetadata(
* to go through the checkout process.
*/
export const createProPlanCheckoutSession = async ({
auth,
owner,
user,
billingPeriod,
planCode,
}: {
auth: Authenticator;
owner: WorkspaceType;
user: UserType;
billingPeriod: BillingPeriod;
planCode: string;
}): Promise<string | null> => {
const stripe = getStripeClient();

const owner = auth.workspace();
if (!owner) {
throw new Error("No workspace found");
}
const user = auth.user();
if (!user) {
throw new Error("No user found");
}

const plan = await Plan.findOne({ where: { code: planCode } });
if (!plan) {
throw new Error(
Expand Down
190 changes: 84 additions & 106 deletions front/lib/resources/subscription_resource.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import type {
Result,
SubscriptionPerSeatPricing,
SubscriptionType,
UserType,
WorkspaceType,
} from "@dust-tt/types";
import { Ok, sendUserOperationMessage } from "@dust-tt/types";
Expand All @@ -20,6 +21,7 @@ import type {
import type Stripe from "stripe";

import { sendProactiveTrialCancelledEmail } from "@app/lib/api/email";
import { getWorkspaceInfos } from "@app/lib/api/workspace";
import type { Authenticator } from "@app/lib/auth";
import { Subscription } from "@app/lib/models/plan";
import { Plan } from "@app/lib/models/plan";
Expand Down Expand Up @@ -145,7 +147,7 @@ export class SubscriptionResource extends BaseResource<Subscription> {
subscriptionResourceByWorkspaceSid[sId] = new SubscriptionResource(
Subscription,
activeSubscription?.get() ||
this.createFreeNoPlanSubscription(workspace.id),
this.createFreeNoPlanSubscription(workspace),
renderPlanFromModel({ plan })
);
}
Expand All @@ -156,10 +158,7 @@ export class SubscriptionResource extends BaseResource<Subscription> {
static async fetchByAuthenticator(
auth: Authenticator
): Promise<SubscriptionResource[]> {
const owner = auth.workspace();
if (!owner) {
throw new Error("Cannot find workspace.");
}
const owner = auth.getNonNullableWorkspace();

const subscriptions = await Subscription.findAll({
where: { workspaceId: owner.id },
Expand All @@ -170,7 +169,7 @@ export class SubscriptionResource extends BaseResource<Subscription> {
(s) =>
new SubscriptionResource(
Subscription,
s,
s.get(),
renderPlanFromModel({ plan: s.plan })
)
);
Expand All @@ -190,7 +189,7 @@ export class SubscriptionResource extends BaseResource<Subscription> {

return new SubscriptionResource(
Subscription,
res,
res.get(),
renderPlanFromModel({ plan: res.plan })
);
}
Expand All @@ -209,11 +208,11 @@ export class SubscriptionResource extends BaseResource<Subscription> {
}): Promise<SubscriptionResource> {
const workspace = await this.findWorkspaceOrThrow(workspaceId);

await this.endActiveSubscription(workspace.id);
await this.endActiveSubscription(workspace);

return new SubscriptionResource(
Subscription,
this.createFreeNoPlanSubscription(workspace.id),
this.createFreeNoPlanSubscription(workspace),
renderPlanFromModel({ plan: FREE_NO_PLAN_DATA })
);
}
Expand Down Expand Up @@ -313,7 +312,7 @@ export class SubscriptionResource extends BaseResource<Subscription> {

return new SubscriptionResource(
Subscription,
newSubscription,
newSubscription.get(),
renderPlanFromModel({ plan: newPlan })
);
}
Expand All @@ -322,10 +321,7 @@ export class SubscriptionResource extends BaseResource<Subscription> {
auth: Authenticator,
enterpriseDetails: EnterpriseUpgradeFormType
) {
const owner = auth.workspace();
if (!owner) {
throw new Error("Cannot find workspace.");
}
const owner = auth.getNonNullableWorkspace();

if (!auth.isDustSuperUser()) {
throw new Error("Cannot upgrade workspace to plan: not allowed.");
Expand Down Expand Up @@ -355,10 +351,7 @@ export class SubscriptionResource extends BaseResource<Subscription> {
auth: Authenticator,
planCode: string
) {
const owner = auth.workspace();
if (!owner) {
throw new Error("Cannot find workspace}");
}
const owner = auth.getNonNullableWorkspace();

if (!auth.isDustSuperUser()) {
throw new Error("Cannot upgrade workspace to plan: not allowed.");
Expand All @@ -374,7 +367,7 @@ export class SubscriptionResource extends BaseResource<Subscription> {
}

// We search for an active subscription for this workspace
const activeSubscription = auth.subscription();
const activeSubscription = auth.subscriptionResource();
if (activeSubscription && activeSubscription.plan.code === newPlan.code) {
throw new Error(
`Cannot subscribe to plan ${planCode}: already subscribed.`
Expand All @@ -401,10 +394,8 @@ export class SubscriptionResource extends BaseResource<Subscription> {
);
}

const isAlreadyOnProPlan = await this.isSubscriptionOnProPlan(
owner,
activeSubscription
);
const isAlreadyOnProPlan =
await activeSubscription.isSubscriptionOnProPlan(owner);

if (!isAlreadyOnProPlan) {
throw new Error(
Expand All @@ -429,63 +420,6 @@ export class SubscriptionResource extends BaseResource<Subscription> {
});
}

static async getCheckoutUrlForUpgrade(
auth: Authenticator,
billingPeriod: BillingPeriod
): Promise<CheckoutUrlResult> {
const owner = auth.workspace();

if (!owner) {
throw new Error(
"Unauthorized `auth` data: cannot process to subscription of new Plan."
);
}

const planCode = owner.metadata?.isBusiness
? PRO_PLAN_SEAT_39_CODE
: PRO_PLAN_SEAT_29_CODE;

const proPlan = await Plan.findOne({
where: { code: PRO_PLAN_SEAT_29_CODE },
});
if (!proPlan) {
throw new Error(`Cannot subscribe to plan ${planCode}: not found.`);
}

const existingSubscription = auth.subscription();

// We verify that the workspace is not already subscribed to the Pro plan product.
if (existingSubscription) {
const isAlreadyOnProPlan = await this.isSubscriptionOnProPlan(
owner,
existingSubscription
);
if (isAlreadyOnProPlan) {
throw new Error(
`Cannot subscribe to plan ${planCode}: already subscribed to a Pro plan.`
);
}
}

// We enter Stripe Checkout flow.
const checkoutUrl = await createProPlanCheckoutSession({
auth,
billingPeriod,
planCode,
});

if (!checkoutUrl) {
throw new Error(
`Cannot subscribe to plan ${planCode}: error while creating Stripe Checkout session (URL is null).`
);
}

return {
checkoutUrl,
plan: renderPlanFromModel({ plan: proPlan }),
};
}

static async maybeCancelInactiveTrials(
auth: Authenticator,
eventStripeSubscription: Stripe.Subscription
Expand Down Expand Up @@ -532,7 +466,7 @@ export class SubscriptionResource extends BaseResource<Subscription> {
const firstAdmin = await getWorkspaceFirstAdmin(workspace);
if (!firstAdmin) {
logger.info(
{ action: "cancelling-trial", workspaceId: workspace.sId },
{ action: "cancelling-trial", workspaceId: auth.workspace()?.sId },
"No first adming found -- skipping email."
);

Expand All @@ -548,6 +482,50 @@ export class SubscriptionResource extends BaseResource<Subscription> {
}
}

async getCheckoutUrlForUpgrade(
owner: WorkspaceType,
user: UserType,
billingPeriod: BillingPeriod
): Promise<CheckoutUrlResult> {
const planCode = owner.metadata?.isBusiness
? PRO_PLAN_SEAT_39_CODE
: PRO_PLAN_SEAT_29_CODE;

const proPlan = await Plan.findOne({
where: { code: PRO_PLAN_SEAT_29_CODE },
});
if (!proPlan) {
throw new Error(`Cannot subscribe to plan ${planCode}: not found.`);
}

// We verify that the workspace is not already subscribed to the Pro plan product.
const isAlreadyOnProPlan = await this.isSubscriptionOnProPlan(owner);
if (isAlreadyOnProPlan) {
throw new Error(
`Cannot subscribe to plan ${planCode}: already subscribed to a Pro plan.`
);
}

// We enter Stripe Checkout flow.
const checkoutUrl = await createProPlanCheckoutSession({
owner,
user,
billingPeriod,
planCode,
});

if (!checkoutUrl) {
throw new Error(
`Cannot subscribe to plan ${planCode}: error while creating Stripe Checkout session (URL is null).`
);
}

return {
checkoutUrl,
plan: renderPlanFromModel({ plan: proPlan }),
};
}

async delete(
auth: Authenticator,
{ transaction }: { transaction?: Transaction } = {}
Expand Down Expand Up @@ -630,14 +608,14 @@ export class SubscriptionResource extends BaseResource<Subscription> {
}

private static createFreeNoPlanSubscription(
workspaceId: number
workspace: LightWorkspaceType
): Attributes<Subscription> {
const now = new Date();
return {
id: FREE_NO_PLAN_SUBSCRIPTION_ID,
sId: generateRandomModelSId(),
status: "ended",
workspaceId: workspaceId,
workspaceId: workspace.id,
createdAt: now,
updatedAt: now,
startDate: now,
Expand All @@ -662,29 +640,10 @@ export class SubscriptionResource extends BaseResource<Subscription> {
);
}

private static async isSubscriptionOnProPlan(
owner: WorkspaceType,
subscription: SubscriptionType
): Promise<boolean> {
if (!subscription.stripeSubscriptionId) {
return false;
}
const stripeSubscription = await getStripeSubscription(
subscription.stripeSubscriptionId
);
if (!stripeSubscription) {
return false;
}

return this.isStripeSubscriptionOnProPlan(owner, stripeSubscription);
}

private static async findWorkspaceOrThrow(
workspaceId: string
): Promise<Workspace> {
const workspace = await Workspace.findOne({
where: { sId: workspaceId },
});
): Promise<LightWorkspaceType> {
const workspace = await getWorkspaceInfos(workspaceId);

if (!workspace) {
throw new Error(`Cannot find workspace ${workspaceId}`);
Expand All @@ -699,13 +658,13 @@ export class SubscriptionResource extends BaseResource<Subscription> {
* @returns The active subscription that was ended, or null if none existed
*/
private static async endActiveSubscription(
workspaceId: number
workspace: LightWorkspaceType
): Promise<Subscription | null> {
const now = new Date();

// Find active subscription
const activeSubscription = await Subscription.findOne({
where: { workspaceId, status: "active" },
where: { workspaceId: workspace.id, status: "active" },
});

if (activeSubscription) {
Expand Down Expand Up @@ -734,4 +693,23 @@ export class SubscriptionResource extends BaseResource<Subscription> {

return activeSubscription;
}

private async isSubscriptionOnProPlan(
owner: WorkspaceType
): Promise<boolean> {
if (!this.stripeSubscriptionId) {
return false;
}
const stripeSubscription = await getStripeSubscription(
this.stripeSubscriptionId
);
if (!stripeSubscription) {
return false;
}

return SubscriptionResource.isStripeSubscriptionOnProPlan(
owner,
stripeSubscription
);
}
}
Loading

0 comments on commit 24dc499

Please sign in to comment.