(core) Billing for formula assistant

Summary:
Adding limits for AI calls and connecting those limits with a Stripe Account.

- New table in homedb called `limits`
- All calls to the AI are not routed through DocApi and measured.
- All products now contain a special key `assistantLimit`, with a default value 0
- Limit is reset every time the subscription has changed its period
- The billing page is updated with two new options that describe the AI plan
- There is a new popup that allows the user to upgrade to a higher plan
- Tiers are read directly from the Stripe product with a volume pricing model

Test Plan: Updated and added

Reviewers: georgegevoian, paulfitz

Reviewed By: georgegevoian

Subscribers: dsagal

Differential Revision: https://phab.getgrist.com/D3907

app/gen-server/entity/BillingAccount.ts

@@ -3,12 +3,13 @@ import {BillingAccountManager} from 'app/gen-server/entity/BillingAccountManager
 import {Organization} from 'app/gen-server/entity/Organization';
 import {Product} from 'app/gen-server/entity/Product';
 import {nativeValues} from 'app/gen-server/lib/values';
+import {Limit} from 'app/gen-server/entity/Limit';
 
 // This type is for billing account status information.  Intended for stuff
 // like "free trial running out in N days".
-interface BillingAccountStatus {
+export interface BillingAccountStatus {
   stripeStatus?: string;
-  currentPeriodEnd?: Date;
+  currentPeriodEnd?: string;
   message?: string;
 }
 
@@ -68,6 +69,9 @@ export class BillingAccount extends BaseEntity {
   @OneToMany(type => Organization, org => org.billingAccount)
   public orgs: Organization[];
 
+  @OneToMany(type => Limit, limit => limit.billingAccount)
+  public limits: Limit[];
+
   // A calculated column that is true if it looks like there is a paid plan.
   @Column({name: 'paid', type: 'boolean', insert: false, select: false})
   public paid?: boolean;

app/gen-server/entity/Limit.ts

@@ -0,0 +1,46 @@
+import {BaseEntity, Column, Entity, JoinColumn, ManyToOne, PrimaryGeneratedColumn} from 'typeorm';
+import {BillingAccount} from 'app/gen-server/entity/BillingAccount';
+import {nativeValues} from 'app/gen-server/lib/values';
+
+@Entity('limits')
+export class Limit extends BaseEntity {
+  @PrimaryGeneratedColumn()
+  public id: number;
+
+  @Column()
+  public limit: number;
+
+  @Column()
+  public usage: number;
+
+  @Column()
+  public type: string;
+
+  @Column({name: 'billing_account_id'})
+  public billingAccountId: number;
+
+  @ManyToOne(type => BillingAccount)
+  @JoinColumn({name: 'billing_account_id'})
+  public billingAccount: BillingAccount;
+
+  @Column({name: 'created_at', default: () => "CURRENT_TIMESTAMP"})
+  public createdAt: Date;
+
+  /**
+   * Last time the Limit.limit value was changed, by an upgrade or downgrade. Null if it has never been changed.
+   */
+  @Column({name: 'changed_at', type: nativeValues.dateTimeType, nullable: true})
+  public changedAt: Date|null;
+
+  /**
+   * Last time the Limit.usage was used (by sending a request to the model). Null if it has never been used.
+   */
+  @Column({name: 'used_at', type: nativeValues.dateTimeType, nullable: true})
+  public usedAt: Date|null;
+
+  /**
+   * Last time the Limit.usage was reset, probably by billing cycle change. Null if it has never been reset.
+   */
+  @Column({name: 'reset_at', type: nativeValues.dateTimeType, nullable: true})
+  public resetAt: Date|null;
+}

app/gen-server/entity/Product.ts

@@ -13,7 +13,11 @@ export const personalLegacyFeatures: Features = {
   // no vanity domain
   maxDocsPerOrg: 10,
   maxSharesPerDoc: 2,
-  maxWorkspacesPerOrg: 1
+  maxWorkspacesPerOrg: 1,
+  /**
+   * One time limit of 100 requests.
+   */
+  baseMaxAssistantCalls: 100,
 };
 
 /**
@@ -23,7 +27,12 @@ export const teamFeatures: Features = {
   workspaces: true,
   vanityDomain: true,
   maxSharesPerWorkspace: 0,   // all workspace shares need to be org members.
-  maxSharesPerDoc: 2
+  maxSharesPerDoc: 2,
+  /**
+   * Limit of 100 requests, but unlike for personal/free orgs the usage for this limit is reset at every billing cycle
+   * through Stripe webhook. For canceled subscription the usage is not reset, as the billing cycle is not changed.
+   */
+  baseMaxAssistantCalls: 100,
 };
 
 /**
@@ -40,6 +49,10 @@ export const teamFreeFeatures: Features = {
   baseMaxDataSizePerDocument: 5000 * 2 * 1024,  // 2KB per row
   baseMaxAttachmentsBytesPerDocument: 1 * 1024 * 1024 * 1024,  // 1GB
   gracePeriodDays: 14,
+  /**
+   * One time limit of 100 requests.
+   */
+  baseMaxAssistantCalls: 100,
 };
 
 /**
@@ -55,6 +68,7 @@ export const teamFreeFeatures: Features = {
   baseMaxDataSizePerDocument: 5000 * 2 * 1024,  // 2KB per row
   baseMaxAttachmentsBytesPerDocument: 1 * 1024 * 1024 * 1024,  // 1GB
   gracePeriodDays: 14,
+  baseMaxAssistantCalls: 100,
 };
 
 export const testDailyApiLimitFeatures = {
@@ -79,6 +93,7 @@ export const suspendedFeatures: Features = {
   maxDocsPerOrg: 0,
   maxSharesPerDoc: 0,
   maxWorkspacesPerOrg: 0,
+  baseMaxAssistantCalls: 0,
 };
 
 /**

app/gen-server/lib/DocApiForwarder.ts

@@ -60,6 +60,7 @@ export class DocApiForwarder {
     app.use('/api/docs/:docId/assign', withDocWithoutAuth);
     app.use('/api/docs/:docId/webhooks/queue', withDoc);
     app.use('/api/docs/:docId/webhooks', withDoc);
+    app.use('/api/docs/:docId/assistant', withDoc);
     app.use('^/api/docs$', withoutDoc);
   }
 

app/gen-server/lib/HomeDBManager.ts

@@ -1,4 +1,4 @@
-import {ApiError} from 'app/common/ApiError';
+import {ApiError, ApiErrorDetails, LimitType} from 'app/common/ApiError';
 import {mapGetOrSet, mapSetOrClear, MapWithTTL} from 'app/common/AsyncCreate';
 import {getDataLimitStatus} from 'app/common/DocLimits';
 import {createEmptyOrgUsageSummary, DocumentUsage, OrgUsageSummary} from 'app/common/DocUsage';
@@ -38,6 +38,7 @@ import {getDefaultProductNames, personalFreeFeatures, Product} from "app/gen-ser
 import {Secret} from "app/gen-server/entity/Secret";
 import {User} from "app/gen-server/entity/User";
 import {Workspace} from "app/gen-server/entity/Workspace";
+import {Limit} from 'app/gen-server/entity/Limit';
 import {Permissions} from 'app/gen-server/lib/Permissions';
 import {scrubUserFromOrg} from "app/gen-server/lib/scrubUserFromOrg";
 import {applyPatch} from 'app/gen-server/lib/TypeORMPatches';
@@ -2880,6 +2881,144 @@ export class HomeDBManager extends EventEmitter {
     return this._org(scope, scope.includeSupport || false, org, options);
   }
 
+  public async getLimits(accountId: number): Promise<Limit[]> {
+    const result = this._connection.transaction(async manager => {
+      return await manager.createQueryBuilder()
+        .select('limit')
+        .from(Limit, 'limit')
+        .innerJoin('limit.billingAccount', 'account')
+        .where('account.id = :accountId', {accountId})
+        .getMany();
+    });
+    return result;
+  }
+
+  public async getLimit(accountId: number, limitType: LimitType): Promise<Limit|null> {
+    return await this._getOrCreateLimit(accountId, limitType, true);
+  }
+
+  public async peekLimit(accountId: number, limitType: LimitType): Promise<Limit|null> {
+    return await this._getOrCreateLimit(accountId, limitType, false);
+  }
+
+  public async removeLimit(scope: Scope, limitType: LimitType): Promise<void> {
+    await this._connection.transaction(async manager => {
+      const org = await this._org(scope, false, scope.org ?? null, {manager, needRealOrg: true})
+        .innerJoinAndSelect('orgs.billingAccount', 'billing_account')
+        .innerJoinAndSelect('billing_account.product', 'product')
+        .leftJoinAndSelect('billing_account.limits', 'limit', 'limit.type = :limitType', {limitType})
+        .getOne();
+      const existing = org?.billingAccount?.limits?.[0];
+      if (existing) {
+        await manager.remove(existing);
+      }
+    });
+  }
+
+  /**
+   * Increases the usage of a limit for a given org. If the limit doesn't exist, it will be created.
+   * Pass `dryRun: true` to check if the limit can be increased without actually increasing it.
+   */
+  public async increaseUsage(scope: Scope, limitType: LimitType, options: {
+    delta: number,
+    dryRun?: boolean,
+  }): Promise<void> {
+    const limitError = await this._connection.transaction(async manager => {
+      const org = await this._org(scope, false, scope.org ?? null, {manager, needRealOrg: true})
+        .innerJoinAndSelect('orgs.billingAccount', 'billing_account')
+        .innerJoinAndSelect('billing_account.product', 'product')
+        .leftJoinAndSelect('billing_account.limits', 'limit', 'limit.type = :limitType', {limitType})
+        .getOne();
+      // If the org doesn't exists, or is a fake one (like for anonymous users), don't do anything.
+      if (!org || org.id === 0) {
+        // This API shouldn't be called, it should be checked first if the org is valid.
+        throw new ApiError(`Can't create a limit for non-existing organization`, 500);
+      }
+      let existing = org?.billingAccount?.limits?.[0];
+      if (!existing) {
+        const product = org?.billingAccount?.product;
+        if (!product) {
+          throw new ApiError(`getLimit: no product found for org`, 500);
+        }
+        if (product.features.baseMaxAssistantCalls === undefined) {
+          // If the product has no assistantLimit, then it is not billable yet, and we don't need to
+          // track usage as it is basically unlimited.
+          return;
+        }
+        existing = new Limit();
+        existing.billingAccountId = org.billingAccountId;
+        existing.type = limitType;
+        existing.limit = product.features.baseMaxAssistantCalls ?? 0;
+        existing.usage = 0;
+      }
+      const limitLess = existing.limit === -1; // -1 means no limit, it is not possible to do in stripe.
+      const usageAfter = existing.usage + options.delta;
+      if (!limitLess && usageAfter > existing.limit) {
+        const billable = Boolean(org?.billingAccount?.stripeCustomerId);
+        return {
+          limit: {
+            maximum: existing.limit,
+            projectedValue: existing.usage + options.delta,
+            quantity: limitType,
+            value: existing.usage,
+          },
+          tips: [{
+            // For non-billable accounts, suggest getting a plan, otherwise suggest visiting the billing page.
+            action: billable ? 'manage' : 'upgrade',
+            message: `Upgrade to a paid plan to increase your ${limitType} limit.`,
+          }]
+        } as ApiErrorDetails;
+      }
+      existing.usage += options.delta;
+      existing.usedAt = new Date();
+      if (!options.dryRun) {
+        await manager.save(existing);
+      }
+    });
+    if (limitError) {
+      let message = `Your ${limitType} limit has been reached. Please upgrade your plan to increase your limit.`;
+      if (limitType === 'assistant') {
+        message = 'You used all available credits. For a bigger limit upgrade you Assistant plan.';
+      }
+      throw new ApiError(message, 429, limitError);
+    }
+  }
+
+  private async _getOrCreateLimit(accountId: number, limitType: LimitType, force: boolean): Promise<Limit|null> {
+    if (accountId === 0) {
+      throw new Error(`getLimit: called for not existing account`);
+    }
+    const result = this._connection.transaction(async manager => {
+      let existing = await manager.createQueryBuilder()
+        .select('limit')
+        .from(Limit, 'limit')
+        .innerJoin('limit.billingAccount', 'account')
+        .where('account.id = :accountId', {accountId})
+          .andWhere('limit.type = :limitType', {limitType})
+        .getOne();
+      if (!force && !existing) { return null; }
+      if (existing) { return existing; }
+      const product = await manager.createQueryBuilder()
+        .select('product')
+        .from(Product, 'product')
+        .innerJoinAndSelect('product.accounts', 'account')
+        .where('account.id = :accountId', {accountId})
+        .getOne();
+      if (!product) {
+        throw new Error(`getLimit: no product for account ${accountId}`);
+      }
+      existing = new Limit();
+      existing.billingAccountId = product.accounts[0].id;
+      existing.type = limitType;
+      existing.limit = product.features.baseMaxAssistantCalls ?? 0;
+      existing.usage = 0;
+      await manager.save(existing);
+      return existing;
+    });
+    return result;
+  }
+
+
   private _org(scope: Scope|null, includeSupport: boolean, org: string|number|null,
                options: QueryOptions = {}): SelectQueryBuilder<Organization> {
     let query = this._orgs(options.manager);

app/gen-server/migration/1685343047786-AssistantLimit.ts

@@ -0,0 +1,82 @@
+import * as sqlUtils from "app/gen-server/sqlUtils";
+import {MigrationInterface, QueryRunner, Table, TableIndex} from 'typeorm';
+
+export class AssistantLimit1685343047786 implements MigrationInterface {
+  public async up(queryRunner: QueryRunner): Promise<void> {
+    const dbType = queryRunner.connection.driver.options.type;
+    const datetime = sqlUtils.datetime(dbType);
+    const now = sqlUtils.now(dbType);
+    await queryRunner.createTable(
+      new Table({
+        name: 'limits',
+        columns: [
+          {
+            name: 'id',
+            type: 'integer',
+            isPrimary: true,
+            isGenerated: true,
+            generationStrategy: 'increment',
+          },
+          {
+            name: 'type',
+            type: 'varchar',
+          },
+          {
+            name: 'billing_account_id',
+            type: 'integer',
+          },
+          {
+            name: 'limit',
+            type: 'integer',
+            default: 0,
+          },
+          {
+            name: 'usage',
+            type: 'integer',
+            default: 0,
+          },
+          {
+            name: "created_at",
+            type: datetime,
+            default: now
+          },
+          {
+            name: "changed_at", // When the limit was last changed
+            type: datetime,
+            isNullable: true
+          },
+          {
+            name: "used_at", // When the usage was last increased
+            type: datetime,
+            isNullable: true
+          },
+          {
+            name: "reset_at", // When the usage was last reset.
+            type: datetime,
+            isNullable: true
+          },
+        ],
+        foreignKeys: [
+          {
+            columnNames: ['billing_account_id'],
+            referencedTableName: 'billing_accounts',
+            referencedColumnNames: ['id'],
+            onDelete: 'CASCADE',
+          },
+        ],
+      })
+    );
+
+    await queryRunner.createIndex(
+      'limits',
+      new TableIndex({
+        name: 'limits_billing_account_id',
+        columnNames: ['billing_account_id'],
+      })
+    );
+  }
+
+  public async down(queryRunner: QueryRunner): Promise<void> {
+    await queryRunner.dropTable('limits');
+  }
+}