feat: Add per user storage quota

3 files changed, 41 insertions, 32 deletions

diff --git a/packages/shared/assetdb.ts b/packages/shared/assetdb.ts
index 77050406..7653ca15 100644
--- a/packages/shared/assetdb.ts
+++ b/packages/shared/assetdb.ts
@@ -16,6 +16,7 @@ import { z } from "zod";
 
 import serverConfig from "./config";
 import logger from "./logger";
+import { QuotaApproved } from "./storageQuota";
 
 const ROOT_PATH = serverConfig.assetsDir;
 
@@ -616,12 +617,22 @@ export async function saveAsset({
   assetId,
   asset,
   metadata,
+  quotaApproved,
 }: {
   userId: string;
   assetId: string;
   asset: Buffer;
   metadata: z.infer<typeof zAssetMetadataSchema>;
+  quotaApproved: QuotaApproved;
 }) {
+  // Verify the quota approval is for the correct user and size
+  if (quotaApproved.userId !== userId) {
+    throw new Error("Quota approval is for a different user");
+  }
+  if (quotaApproved.approvedSize < asset.byteLength) {
+    throw new Error("Asset size exceeds approved quota");
+  }
+
   return defaultAssetStore.saveAsset({ userId, assetId, asset, metadata });
 }
 
@@ -630,12 +641,22 @@ export async function saveAssetFromFile({
   assetId,
   assetPath,
   metadata,
+  quotaApproved,
 }: {
   userId: string;
   assetId: string;
   assetPath: string;
   metadata: z.infer<typeof zAssetMetadataSchema>;
+  quotaApproved: QuotaApproved;
 }) {
+  // Verify the quota approval is for the correct user
+  if (quotaApproved.userId !== userId) {
+    throw new Error("Quota approval is for a different user");
+  }
+
+  // For file-based saves, we'll verify the file size matches the approved size
+  // when the underlying store implementation reads the file
+
   return defaultAssetStore.saveAssetFromFile({
     userId,
     assetId,
@@ -724,35 +745,3 @@ export async function deleteUserAssets({ userId }: { userId: string }) {
 export async function* getAllAssets() {
   yield* defaultAssetStore.getAllAssets();
 }
-
-export async function storeScreenshot(
-  screenshot: Buffer | undefined,
-  userId: string,
-  jobId: string,
-) {
-  if (!serverConfig.crawler.storeScreenshot) {
-    logger.info(
-      `[Crawler][${jobId}] Skipping storing the screenshot as per the config.`,
-    );
-    return null;
-  }
-  if (!screenshot) {
-    logger.info(
-      `[Crawler][${jobId}] Skipping storing the screenshot as it's empty.`,
-    );
-    return null;
-  }
-  const assetId = newAssetId();
-  const contentType = "image/png";
-  const fileName = "screenshot.png";
-  await saveAsset({
-    userId,
-    assetId,
-    metadata: { contentType, fileName },
-    asset: screenshot,
-  });
-  logger.info(
-    `[Crawler][${jobId}] Stored the screenshot as assetId: ${assetId}`,
-  );
-  return { assetId, contentType, fileName, size: screenshot.byteLength };
-}
diff --git a/packages/shared/storageQuota.ts b/packages/shared/storageQuota.ts
new file mode 100644
index 00000000..6b7441a2
--- /dev/null
+++ b/packages/shared/storageQuota.ts
@@ -0,0 +1,19 @@
+/**
+ * A token that proves storage quota has been checked and approved.
+ * This class cannot be instantiated directly - it can only be created
+ * by the checkStorageQuota function.
+ */
+export class QuotaApproved {
+  private constructor(
+    public readonly userId: string,
+    public readonly approvedSize: number,
+  ) {}
+
+  /**
+   * Internal method to create a QuotaApproved token.
+   * This should only be called by checkStorageQuota.
+   */
+  static _create(userId: string, approvedSize: number): QuotaApproved {
+    return new QuotaApproved(userId, approvedSize);
+  }
+}
diff --git a/packages/shared/types/admin.ts b/packages/shared/types/admin.ts
index 652ee2d3..f4c62af0 100644
--- a/packages/shared/types/admin.ts
+++ b/packages/shared/types/admin.ts
@@ -12,6 +12,7 @@ export const updateUserSchema = z.object({
   userId: z.string(),
   role: z.enum(["user", "admin"]).optional(),
   bookmarkQuota: z.number().int().min(0).nullable().optional(),
+  storageQuota: z.number().int().min(0).nullable().optional(),
 });
 
 export const resetPasswordSchema = z