From b36faeaf540add44f0921bd6ff0d5a4360f5eb5d Mon Sep 17 00:00:00 2001
From: Nicolas <nicolascamara29@gmail.com>
Date: Tue, 20 Aug 2024 14:39:52 -0300
Subject: [PATCH] Nick:

---
 apps/api/src/controllers/v1/crawl.ts |  5 +++++
 apps/api/src/controllers/v1/types.ts | 11 +++++++++--
 apps/api/src/routes/v1.ts            | 11 ++++++++---
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/apps/api/src/controllers/v1/crawl.ts b/apps/api/src/controllers/v1/crawl.ts
index 51055fb5..6166af15 100644
--- a/apps/api/src/controllers/v1/crawl.ts
+++ b/apps/api/src/controllers/v1/crawl.ts
@@ -32,9 +32,14 @@ export async function crawlController(
 
   await logCrawl(id, req.auth.team_id);
 
+  const { remainingCredits } = req.account;
+
+  // TODO: Get rid of crawlerOptions
   const crawlerOptions = legacyCrawlerOptions(req.body.crawlerOptions);
   const pageOptions = legacyScrapeOptions(req.body.scrapeOptions);
 
+  crawlerOptions.limit = Math.min(remainingCredits, crawlerOptions.limit);
+  
   const sc: StoredCrawl = {
     originUrl: req.body.url,
     crawlerOptions,
diff --git a/apps/api/src/controllers/v1/types.ts b/apps/api/src/controllers/v1/types.ts
index 3fe35891..39f34abe 100644
--- a/apps/api/src/controllers/v1/types.ts
+++ b/apps/api/src/controllers/v1/types.ts
@@ -100,9 +100,10 @@ export type CrawlerOptions = z.infer<typeof crawlerOptions>;
 export const crawlRequestSchema = z.object({
   url,
   origin: z.string().optional().default("api"),
-  crawlerOptions: crawlerOptions.default({}),
+  crawlerOptions: crawlerOptions.default({}), // TODO: Get rid of this
   scrapeOptions: scrapeOptions.omit({ timeout: true }).default({}),
   webhook: z.string().url().optional(),
+  limit: z.number().default(10000), //
 });
 
 // export type CrawlRequest = {
@@ -225,20 +226,26 @@ type AuthObject = {
   plan: string;
 };
 
+type Account = {
+  remainingCredits: number;
+};
+
 export interface RequestWithMaybeAuth<
   ReqParams = {},
   ReqBody = undefined,
   ResBody = undefined
 > extends Request<ReqParams, ReqBody, ResBody> {
   auth?: AuthObject;
+  account?: Account;
 }
 
 export interface RequestWithAuth<
   ReqParams = {},
   ReqBody = undefined,
-  ResBody = undefined
+  ResBody = undefined,
 > extends Request<ReqParams, ReqBody, ResBody> {
   auth: AuthObject;
+  account?: Account;
 }
 
 export function legacyCrawlerOptions(x: CrawlerOptions) {
diff --git a/apps/api/src/routes/v1.ts b/apps/api/src/routes/v1.ts
index 25e12c63..51f07c85 100644
--- a/apps/api/src/routes/v1.ts
+++ b/apps/api/src/routes/v1.ts
@@ -24,12 +24,17 @@ import { isUrlBlocked } from "../scraper/WebScraper/utils/blocklist";
 // import { livenessController } from "../controllers/v1/liveness";
 // import { readinessController } from "../controllers/v1/readiness";
 
-function checkCreditsMiddleware(minimum: number): (req: RequestWithAuth, res: Response, next: NextFunction) => void {
+function checkCreditsMiddleware(minimum?: number): (req: RequestWithAuth, res: Response, next: NextFunction) => void {
     return (req, res, next) => {
         (async () => {
-            if (!(await checkTeamCredits(req.auth.team_id, minimum)).success) {
+            if (!minimum && req.body) {
+                minimum = (req.body as any)?.limit ?? 1;
+            }
+            const { success, message, remainingCredits } = await checkTeamCredits(req.auth.team_id, minimum);
+            if (!success) {
                 return res.status(402).json({ success: false, error: "Insufficient credits" });
             }
+            req.account = { remainingCredits }
             next();
         })()
             .catch(err => next(err));
@@ -101,7 +106,7 @@ v1Router.post(
     blocklistMiddleware,
     authMiddleware(RateLimiterMode.Crawl),
     idempotencyMiddleware,
-    checkCreditsMiddleware(1),
+    checkCreditsMiddleware(),
     wrap(crawlController)
 );