Add: script to generate reimbursement and emails for the keys incident

front/migrations/20241101_workspace_keys_incident.ts

@@ -0,0 +1,213 @@
+import { parse } from "csv-parse";
+import * as fs from "fs";
+import { exit } from "process";
+
+import { sendEmail } from "@app/lib/api/email";
+import { getMembers } from "@app/lib/api/workspace";
+import { Authenticator } from "@app/lib/auth";
+import { Workspace } from "@app/lib/models/workspace";
+import { getStripeClient } from "@app/lib/plans/stripe";
+import type { Logger } from "@app/logger/logger";
+import { makeScript } from "@app/scripts/helpers";
+
+const MIN_AMOUNT_FOR_REIMBURSEMENT = 1;
+const CURRENCY = "eur";
+
+const sendIncidentEmailToAdmins = async (
+  workspaceId: number,
+  cost_in_cents: number,
+  reimbursement_in_cents: number,
+  logger: Logger
+) => {
+  const message = {
+    from: {
+      name: "Dust team",
+      email: "[email protected]",
+    },
+    subject: `[Dust] Workspace keys incident`,
+    html: `<p>Hello from Dust,</p>
+      <p>We had an incident on Thuesday 31st of October 2024 between 12:44 ECT and 13:20 ECT that affected your workspace.</p>
+      <p>During that time, as you had set up provider keys for your workspace, they have been used instead of our own keys.</p>
+      <p>In your case, the incident caused a cost of ${cost_in_cents / 100} USD in tokens. We will reimburse you via a ${reimbursement_in_cents / 100} EUR discount on your next invoice.</p>
+      <p>Please reply to this email if you have any questions.</p>
+      <p>The Dust team</p>`,
+  };
+
+  const workspace = await Workspace.findByPk(workspaceId);
+  if (!workspace) {
+    // Should not happen but just in case
+    logger.error(`Workspace ${workspaceId} not found. Skipping email.`);
+    return;
+  }
+
+  const auth = await Authenticator.internalAdminForWorkspace(workspace.sId);
+
+  const { members: admins } = await getMembers(auth, { roles: ["admin"] });
+  const adminEmails = admins.map((u) => u.email);
+
+  return Promise.all(adminEmails.map((email) => sendEmail(email, message)));
+};
+
+// Basic file read
+const readCsvFile = async (filePath: string) => {
+  const parser = parse({
+    delimiter: ",",
+    columns: true, // Use first row as headers
+    skip_empty_lines: true,
+    trim: true,
+    cast: function (value, context) {
+      if (
+        [
+          "workspaceId",
+          "number_of_runs",
+          "prompt_tokens",
+          "completion_tokens",
+        ].includes(context.column as string)
+      ) {
+        // remove the "," from the value and parse as integer
+        return parseInt(value.replace(/,/g, ""));
+      } else {
+        return value;
+      }
+    },
+  });
+
+  const records: any[] = [];
+  fs.createReadStream(filePath).pipe(parser);
+
+  for await (const record of parser) {
+    records.push(record);
+  }
+
+  return records;
+};
+
+makeScript(
+  {
+    csvPath: {
+      alias: "csv",
+      describe: "Path to the CSV file",
+      type: "string",
+    },
+    mergeCoupons: {
+      alias: "merge",
+      describe: "Merge coupons",
+      type: "boolean",
+      default: false,
+    },
+  },
+  async ({ csvPath, mergeCoupons, execute }, logger) => {
+    const records = await readCsvFile(csvPath);
+    logger.info(`Read ${records.length} records from CSV file`);
+    const stripe = getStripeClient();
+
+    // Check that no existing coupons are present
+    logger.info(`Checking existing coupons...`);
+    const alreadyHasCoupons = [];
+    for (const record of records) {
+      if (
+        record["stripe_subscription_id"] &&
+        record["total_tokens_cost_in_dollars"] > MIN_AMOUNT_FOR_REIMBURSEMENT
+      ) {
+        // Check if the customer already has a coupon
+        const subscriptionId = record["stripe_subscription_id"];
+
+        // Check if the customer already has a coupon
+        const subscription =
+          await stripe.subscriptions.retrieve(subscriptionId);
+
+        if (subscription.discount && subscription.discount.coupon.amount_off) {
+          alreadyHasCoupons.push(record);
+        }
+      }
+    }
+
+    if (alreadyHasCoupons.length > 0 && !mergeCoupons) {
+      // Stop the script with an error
+      logger.error(
+        `Some records already have a discount: ${alreadyHasCoupons.map((r) => r["name"]).join(", ")}. Use --merge if you want to merge the existing discount with the new one.`
+      );
+      exit(1);
+    }
+
+    // Process records
+    for (const record of records) {
+      if (
+        record["total_tokens_cost_in_dollars"] > MIN_AMOUNT_FOR_REIMBURSEMENT
+      ) {
+        if (record["stripe_subscription_id"]) {
+          const subscriptionId = record["stripe_subscription_id"];
+
+          const amount =
+            Math.ceil(record["total_tokens_cost_in_dollars"]) * 100;
+          let finalAmount = amount;
+
+          // Check if the customer already has a coupon
+          const subscription =
+            await stripe.subscriptions.retrieve(subscriptionId);
+
+          if (
+            subscription.discount &&
+            subscription.discount.coupon.amount_off
+          ) {
+            if (!mergeCoupons) {
+              // Should not happen as we are checking this before
+              // Just in case: stop the script with an error
+              throw new Error(
+                `Record ${record["name"]} (${record["workspaceId"]}) already has a discount of ${subscription.discount.coupon.amount_off} cents. Use --merge to merge the discounts.`
+              );
+            } else {
+              // Update the amount
+              const currentAmount = subscription.discount.coupon.amount_off;
+              finalAmount += currentAmount;
+
+              logger.info(
+                `Customer already has a coupon: adding ${currentAmount} cents to the amount, now ${amount} cents.`
+              );
+            }
+          }
+
+          if (execute) {
+            // Log
+            logger.info(
+              `Creating one-time coupon for ${record["name"]} (${record["workspaceId"]}) with amount ${amount} cents`
+            );
+
+            // Create a one-time coupon
+            const coupon = await stripe.coupons.create({
+              duration: "once",
+              amount_off: finalAmount, // amount in cents
+              currency: CURRENCY, // here we apply the discount in euros considering a rate of 1:1 with usd
+              name: "One-time discount",
+            });
+
+            // Apply coupon to subscription
+            await stripe.subscriptions.update(subscriptionId, {
+              coupon: coupon.id,
+            });
+
+            // Send email to admins
+            await sendIncidentEmailToAdmins(
+              record["workspaceId"],
+              record["total_tokens_cost_in_dollars"] * 100,
+              amount,
+              logger
+            );
+          } else {
+            logger.info(
+              `Dry run: would have created one-time coupon for ${record["name"]} (${record["workspaceId"]}) with amount ${amount} cents`
+            );
+          }
+        } else {
+          logger.error(
+            `Skipping record ${record["name"]} (${record["workspaceId"]}) without stripe_subscription_id`
+          );
+        }
+      } else {
+        logger.info(
+          `Skipping record ${record["name"]} (${record["workspaceId"]}) with total_tokens_cost_in_dollars less than ${MIN_AMOUNT_FOR_REIMBURSEMENT}$: (${record["total_tokens_cost_in_dollars"]})`
+        );
+      }
+    }
+  }
+);