AI/firecrawl
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/mendableai/firecrawl synced 2025-08-14 18:36:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

ACUC: Dynamic Limits (FIR-1641) (#1434)

Browse Source 
* extend acuc definition

* kill plan

* stuff

* stupid tests

* feat: better acuc

* feat(acuc): mock ACUC when not using db auth
This commit is contained in:
Gergő Móricz 2025-04-10 18:49:23 +02:00 committed by GitHub
parent f2865f6699
commit 6a10f0689d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
43 changed files with 611 additions and 985 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
apps/api/src/__tests__/concurrency-limit.test.ts
View File

@ -9,8 +9,6 @@ import {
getConcurrencyQueueJobsCount, getConcurrencyQueueJobsCount,
ConcurrencyLimitedJob, ConcurrencyLimitedJob,
} from "../lib/concurrency-limit"; } from "../lib/concurrency-limit";
import { CONCURRENCY_LIMIT, getConcurrencyLimitMax } from "../services/rate-limiter";
import { PlanType } from "../types";
// Mock Redis client // Mock Redis client
jest.mock("../services/queue-service", () => ({ jest.mock("../services/queue-service", () => ({
@ -174,34 +172,6 @@ describe("Concurrency Limit", () => {
}); });
}); });
describe("getConcurrencyLimitMax", () => {
it("should return correct limit for free plan", () => {
const result = getConcurrencyLimitMax("free");
expect(result).toBe(2);
});
it("should return correct limit for standard plan", () => {
const result = getConcurrencyLimitMax("standard");
expect(result).toBe(CONCURRENCY_LIMIT.standard);
});
it("should return correct limit for scale plan", () => {
const result = getConcurrencyLimitMax("scale");
expect(result).toBe(CONCURRENCY_LIMIT.scale);
});
it("should return default limit for unknown plan", () => {
const result = getConcurrencyLimitMax("unknown" as PlanType);
expect(result).toBe(10);
});
it("should handle special team IDs", () => {
process.env.DEV_B_TEAM_ID = "dev-b-team";
const result = getConcurrencyLimitMax("free", "dev-b-team");
expect(result).toBe(120);
});
});
describe("Integration Scenarios", () => { describe("Integration Scenarios", () => {
it("should handle complete job lifecycle", async () => { it("should handle complete job lifecycle", async () => {
const mockJob: ConcurrencyLimitedJob = { const mockJob: ConcurrencyLimitedJob = {

1
apps/api/src/__tests__/deep-research/unit/deep-research-redis.test.ts
View File

@ -20,7 +20,6 @@ describe("Deep Research Redis Operations", () => {
const mockResearch: StoredDeepResearch = { const mockResearch: StoredDeepResearch = {
id: "test-id", id: "test-id",
team_id: "team-1", team_id: "team-1",
plan: "pro",
createdAt: Date.now(), createdAt: Date.now(),
status: "processing", status: "processing",
currentDepth: 0, currentDepth: 0,

20
apps/api/src/__tests__/queue-concurrency-integration.test.ts
View File

@ -6,8 +6,8 @@ import {
takeConcurrencyLimitedJob, takeConcurrencyLimitedJob,
removeConcurrencyLimitActiveJob, removeConcurrencyLimitActiveJob,
} from "../lib/concurrency-limit"; } from "../lib/concurrency-limit";
import { getConcurrencyLimitMax } from "../services/rate-limiter"; import { WebScraperOptions } from "../types";
import { WebScraperOptions, PlanType } from "../types"; import { getACUCTeam } from "../controllers/auth";
// Mock all the dependencies // Mock all the dependencies
const mockAdd = jest.fn(); const mockAdd = jest.fn();
@ -32,7 +32,6 @@ jest.mock("uuid", () => ({
describe("Queue Concurrency Integration", () => { describe("Queue Concurrency Integration", () => {
const mockTeamId = "test-team-id"; const mockTeamId = "test-team-id";
const mockPlan = "standard" as PlanType;
const mockNow = Date.now(); const mockNow = Date.now();
const defaultScrapeOptions = { const defaultScrapeOptions = {
@ -77,7 +76,6 @@ describe("Queue Concurrency Integration", () => {
url: "https://test.com", url: "https://test.com",
mode: "single_urls", mode: "single_urls",
team_id: mockTeamId, team_id: mockTeamId,
plan: mockPlan,
scrapeOptions: defaultScrapeOptions, scrapeOptions: defaultScrapeOptions,
crawlerOptions: null, crawlerOptions: null,
}; };
@ -104,8 +102,10 @@ describe("Queue Concurrency Integration", () => {
it("should add job to concurrency queue when at concurrency limit", async () => { it("should add job to concurrency queue when at concurrency limit", async () => {
// Mock current active jobs to be at limit // Mock current active jobs to be at limit
const maxConcurrency = getConcurrencyLimitMax(mockPlan); (getACUCTeam as jest.Mock).mockResolvedValue({
const activeJobs = Array(maxConcurrency).fill("active-job"); concurrency: 15,
} as any);
const activeJobs = Array(15).fill("active-job");
(redisConnection.zrangebyscore as jest.Mock).mockResolvedValue( (redisConnection.zrangebyscore as jest.Mock).mockResolvedValue(
activeJobs, activeJobs,
); );
@ -136,7 +136,6 @@ describe("Queue Concurrency Integration", () => {
url: `https://test${i}.com`, url: `https://test${i}.com`,
mode: "single_urls", mode: "single_urls",
team_id: mockTeamId, team_id: mockTeamId,
plan: mockPlan,
scrapeOptions: defaultScrapeOptions, scrapeOptions: defaultScrapeOptions,
} as WebScraperOptions, } as WebScraperOptions,
opts: { opts: {
@ -146,7 +145,10 @@ describe("Queue Concurrency Integration", () => {
})); }));
it("should handle batch jobs respecting concurrency limits", async () => { it("should handle batch jobs respecting concurrency limits", async () => {
const maxConcurrency = getConcurrencyLimitMax(mockPlan); const maxConcurrency = 15;
(getACUCTeam as jest.Mock).mockResolvedValue({
concurrency: maxConcurrency,
} as any);
const totalJobs = maxConcurrency + 5; // Some jobs should go to queue const totalJobs = maxConcurrency + 5; // Some jobs should go to queue
const mockJobs = createMockJobs(totalJobs); const mockJobs = createMockJobs(totalJobs);
@ -180,7 +182,6 @@ describe("Queue Concurrency Integration", () => {
id: "test-job", id: "test-job",
data: { data: {
team_id: mockTeamId, team_id: mockTeamId,
plan: mockPlan,
}, },
}; };
@ -218,7 +219,6 @@ describe("Queue Concurrency Integration", () => {
id: "failing-job", id: "failing-job",
data: { data: {
team_id: mockTeamId, team_id: mockTeamId,
plan: mockPlan,
}, },
}; };

4
apps/api/src/__tests__/snips/crawl.test.ts
View File

@ -53,7 +53,7 @@ describe("Crawl tests", () => {
// expect(page.metadata.url ?? page.metadata.sourceURL).toMatch(/^https:\/\/(www\.)?firecrawl\.dev\/blog/); // expect(page.metadata.url ?? page.metadata.sourceURL).toMatch(/^https:\/\/(www\.)?firecrawl\.dev\/blog/);
// } // }
// } // }
// }, 120000); // }, 300000);
// TEMP: Flaky // TEMP: Flaky
// it.concurrent("discovers URLs properly when maxDiscoveryDepth is provided", async () => { // it.concurrent("discovers URLs properly when maxDiscoveryDepth is provided", async () => {
@ -71,5 +71,5 @@ describe("Crawl tests", () => {
// expect(page.metadata.url ?? page.metadata.sourceURL).not.toMatch(/^https:\/\/(www\.)?firecrawl\.dev\/blog\/.+$/); // expect(page.metadata.url ?? page.metadata.sourceURL).not.toMatch(/^https:\/\/(www\.)?firecrawl\.dev\/blog\/.+$/);
// } // }
// } // }
// }, 120000); // }, 300000);
}); });

308
apps/api/src/controllers/auth.ts
View File

@ -3,7 +3,6 @@ import { getRateLimiter, isTestSuiteToken } from "../services/rate-limiter";
import { import {
AuthResponse, AuthResponse,
NotificationType, NotificationType,
PlanType,
RateLimiterMode, RateLimiterMode,
} from "../types"; } from "../types";
import { supabase_rr_service, supabase_service } from "../services/supabase"; import { supabase_rr_service, supabase_service } from "../services/supabase";
@ -16,7 +15,7 @@ import { deleteKey, getValue } from "../services/redis";
import { setValue } from "../services/redis"; import { setValue } from "../services/redis";
import { validate } from "uuid"; import { validate } from "uuid";
import * as Sentry from "@sentry/node"; import * as Sentry from "@sentry/node";
import { AuthCreditUsageChunk } from "./v1/types"; import { AuthCreditUsageChunk, AuthCreditUsageChunkFromTeam } from "./v1/types";
// const { data, error } = await supabase_service // const { data, error } = await supabase_service
// .from('api_keys') // .from('api_keys')
// .select(` // .select(`
@ -38,12 +37,13 @@ function normalizedApiIsUuid(potentialUuid: string): boolean {
export async function setCachedACUC( export async function setCachedACUC(
api_key: string, api_key: string,
is_extract: boolean,
acuc: acuc:
| AuthCreditUsageChunk | AuthCreditUsageChunk
| null | null
| ((acuc: AuthCreditUsageChunk) => AuthCreditUsageChunk | null), | ((acuc: AuthCreditUsageChunk) => AuthCreditUsageChunk | null),
) { ) {
const cacheKeyACUC = `acuc_${api_key}`; const cacheKeyACUC = `acuc_${api_key}_${is_extract ? "extract" : "scrape"}`;
const redLockKey = `lock_${cacheKeyACUC}`; const redLockKey = `lock_${cacheKeyACUC}`;
try { try {
@ -72,13 +72,55 @@ export async function setCachedACUC(
} }
} }
const mockACUC: () => AuthCreditUsageChunk = () => ({
api_key: "bypass",
team_id: "bypass",
sub_id: "bypass",
sub_current_period_start: new Date().toISOString(),
sub_current_period_end: new Date(new Date().getTime() + 30 * 24 * 60 * 60 * 1000).toISOString(),
sub_user_id: "bypass",
price_id: "bypass",
rate_limits: {
crawl: 99999999,
scrape: 99999999,
extract: 99999999,
search: 99999999,
map: 99999999,
preview: 99999999,
crawlStatus: 99999999,
extractStatus: 99999999,
},
price_credits: 99999999,
credits_used: 0,
coupon_credits: 99999999,
adjusted_credits_used: 0,
remaining_credits: 99999999,
total_credits_sum: 99999999,
plan_priority: {
bucketLimit: 25,
planModifier: 0.1,
},
concurrency: 99999999,
is_extract: false,
});
export async function getACUC( export async function getACUC(
api_key: string, api_key: string,
cacheOnly = false, cacheOnly = false,
useCache = true, useCache = true,
mode?: RateLimiterMode, mode?: RateLimiterMode,
): Promise<AuthCreditUsageChunk | null> { ): Promise<AuthCreditUsageChunk | null> {
const cacheKeyACUC = `acuc_${api_key}_${mode}`; let isExtract =
mode === RateLimiterMode.Extract ||
mode === RateLimiterMode.ExtractStatus;
if (process.env.USE_DB_AUTHENTICATION !== "true") {
const acuc = mockACUC();
acuc.is_extract = isExtract;
return acuc;
}
const cacheKeyACUC = `acuc_${api_key}_${isExtract ? "extract" : "scrape"}`;
// if (useCache) { // if (useCache) {
// const cachedACUC = await getValue(cacheKeyACUC); // const cachedACUC = await getValue(cacheKeyACUC);
@ -92,15 +134,11 @@ export async function getACUC(
let error; let error;
let retries = 0; let retries = 0;
const maxRetries = 5; const maxRetries = 5;
let isExtract =
mode === RateLimiterMode.Extract ||
mode === RateLimiterMode.ExtractStatus;
while (retries < maxRetries) { while (retries < maxRetries) {
const client = const client =
Math.random() > (2/3) ? supabase_rr_service : supabase_service; Math.random() > (2/3) ? supabase_rr_service : supabase_service;
({ data, error } = await client.rpc( ({ data, error } = await client.rpc(
"auth_credit_usage_chunk_27_tally", "auth_credit_usage_chunk_28",
{ input_key: api_key, i_is_extract: isExtract, tally_untallied_credits: true }, { input_key: api_key, i_is_extract: isExtract, tally_untallied_credits: true },
{ get: true }, { get: true },
)); ));
@ -125,6 +163,117 @@ export async function getACUC(
await new Promise((resolve) => setTimeout(resolve, 200)); await new Promise((resolve) => setTimeout(resolve, 200));
} }
const chunk: AuthCreditUsageChunk | null =
data.length === 0 ? null : data[0].team_id === null ? null : data[0];
// NOTE: Should we cache null chunks? - mogery
// if (chunk !== null && useCache) {
// setCachedACUC(api_key, isExtract, chunk);
// }
return chunk ? { ...chunk, is_extract: isExtract } : null;
// } else {
// return null;
// }
}
export async function setCachedACUCTeam(
team_id: string,
is_extract: boolean,
acuc:
| AuthCreditUsageChunkFromTeam
| null
| ((acuc: AuthCreditUsageChunkFromTeam) => AuthCreditUsageChunkFromTeam | null),
) {
const cacheKeyACUC = `acuc_team_${team_id}_${is_extract ? "extract" : "scrape"}`;
const redLockKey = `lock_${cacheKeyACUC}`;
try {
await redlock.using([redLockKey], 10000, {}, async (signal) => {
if (typeof acuc === "function") {
acuc = acuc(JSON.parse((await getValue(cacheKeyACUC)) ?? "null"));
if (acuc === null) {
if (signal.aborted) {
throw signal.error;
}
return;
}
}
if (signal.aborted) {
throw signal.error;
}
// Cache for 1 hour. - mogery
await setValue(cacheKeyACUC, JSON.stringify(acuc), 3600, true);
});
} catch (error) {
logger.error(`Error updating cached ACUC ${cacheKeyACUC}: ${error}`);
}
}
export async function getACUCTeam(
team_id: string,
cacheOnly = false,
useCache = true,
mode?: RateLimiterMode,
): Promise<AuthCreditUsageChunkFromTeam | null> {
let isExtract =
mode === RateLimiterMode.Extract ||
mode === RateLimiterMode.ExtractStatus;
if (process.env.USE_DB_AUTHENTICATION !== "true") {
const acuc = mockACUC();
acuc.is_extract = isExtract;
return acuc;
}
const cacheKeyACUC = `acuc_team_${team_id}_${isExtract ? "extract" : "scrape"}`;
// if (useCache) {
// const cachedACUC = await getValue(cacheKeyACUC);
// if (cachedACUC !== null) {
// return JSON.parse(cachedACUC);
// }
// }
// if (!cacheOnly) {
let data;
let error;
let retries = 0;
const maxRetries = 5;
while (retries < maxRetries) {
const client =
Math.random() > (2/3) ? supabase_rr_service : supabase_service;
({ data, error } = await client.rpc(
"auth_credit_usage_chunk_28_from_team",
{ input_team: team_id, i_is_extract: isExtract, tally_untallied_credits: true },
{ get: true },
));
if (!error) {
break;
}
logger.warn(
`Failed to retrieve authentication and credit usage data after ${retries}, trying again...`,
{ error }
);
retries++;
if (retries === maxRetries) {
throw new Error(
"Failed to retrieve authentication and credit usage data after 3 attempts: " +
JSON.stringify(error),
);
}
// Wait for a short time before retrying
await new Promise((resolve) => setTimeout(resolve, 200));
}
const chunk: AuthCreditUsageChunk | null = const chunk: AuthCreditUsageChunk | null =
data.length === 0 ? null : data[0].team_id === null ? null : data[0]; data.length === 0 ? null : data[0].team_id === null ? null : data[0];
@ -141,10 +290,10 @@ export async function getACUC(
export async function clearACUC(api_key: string): Promise<void> { export async function clearACUC(api_key: string): Promise<void> {
// Delete cache for all rate limiter modes // Delete cache for all rate limiter modes
const modes = Object.values(RateLimiterMode); const modes = [true, false];
await Promise.all( await Promise.all(
modes.map(async (mode) => { modes.map(async (mode) => {
const cacheKey = `acuc_${api_key}_${mode}`; const cacheKey = `acuc_${api_key}_${mode ? "extract" : "scrape"}`;
await deleteKey(cacheKey); await deleteKey(cacheKey);
}), }),
); );
@ -153,6 +302,20 @@ export async function clearACUC(api_key: string): Promise<void> {
await deleteKey(`acuc_${api_key}`); await deleteKey(`acuc_${api_key}`);
} }
export async function clearACUCTeam(team_id: string): Promise<void> {
// Delete cache for all rate limiter modes
const modes = [true, false];
await Promise.all(
modes.map(async (mode) => {
const cacheKey = `acuc_team_${team_id}_${mode ? "extract" : "scrape"}`;
await deleteKey(cacheKey);
}),
);
// Also clear the base cache key
await deleteKey(`acuc_team_${team_id}`);
}
export async function authenticateUser( export async function authenticateUser(
req, req,
res, res,
@ -192,13 +355,12 @@ export async function supaAuthenticateUser(
const iptoken = incomingIP + token; const iptoken = incomingIP + token;
let rateLimiter: RateLimiterRedis; let rateLimiter: RateLimiterRedis;
let subscriptionData: { team_id: string; plan: string } | null = null; let subscriptionData: { team_id: string} | null = null;
let normalizedApi: string; let normalizedApi: string;
let teamId: string | null = null; let teamId: string | null = null;
let priceId: string | null = null; let priceId: string | null = null;
let chunk: AuthCreditUsageChunk | null = null; let chunk: AuthCreditUsageChunk | null = null;
let plan: PlanType = "free";
if (token == "this_is_just_a_preview_token") { if (token == "this_is_just_a_preview_token") {
throw new Error( throw new Error(
"Unauthenticated Playground calls are temporarily disabled due to abuse. Please sign up.", "Unauthenticated Playground calls are temporarily disabled due to abuse. Please sign up.",
@ -213,7 +375,6 @@ export async function supaAuthenticateUser(
rateLimiter = getRateLimiter(RateLimiterMode.Preview, token); rateLimiter = getRateLimiter(RateLimiterMode.Preview, token);
} }
teamId = `preview_${iptoken}`; teamId = `preview_${iptoken}`;
plan = "free";
} else { } else {
normalizedApi = parseApi(token); normalizedApi = parseApi(token);
if (!normalizedApiIsUuid(normalizedApi)) { if (!normalizedApiIsUuid(normalizedApi)) {
@ -237,65 +398,13 @@ export async function supaAuthenticateUser(
teamId = chunk.team_id; teamId = chunk.team_id;
priceId = chunk.price_id; priceId = chunk.price_id;
plan = getPlanByPriceId(priceId);
subscriptionData = { subscriptionData = {
team_id: teamId, team_id: teamId,
plan,
}; };
switch (mode) { rateLimiter = getRateLimiter(
case RateLimiterMode.Crawl: mode ?? RateLimiterMode.Crawl,
rateLimiter = getRateLimiter( chunk.rate_limits,
RateLimiterMode.Crawl, );
token,
subscriptionData.plan,
);
break;
case RateLimiterMode.Scrape:
rateLimiter = getRateLimiter(
RateLimiterMode.Scrape,
token,
subscriptionData.plan,
teamId,
);
break;
case RateLimiterMode.Search:
rateLimiter = getRateLimiter(
RateLimiterMode.Search,
token,
subscriptionData.plan,
);
break;
case RateLimiterMode.Map:
rateLimiter = getRateLimiter(
RateLimiterMode.Map,
token,
subscriptionData.plan,
);
break;
case RateLimiterMode.Extract:
rateLimiter = getRateLimiter(
RateLimiterMode.Extract,
token,
subscriptionData.plan,
);
break;
case RateLimiterMode.ExtractStatus:
rateLimiter = getRateLimiter(RateLimiterMode.ExtractStatus, token);
break;
case RateLimiterMode.CrawlStatus:
rateLimiter = getRateLimiter(RateLimiterMode.CrawlStatus, token);
break;
case RateLimiterMode.Preview:
rateLimiter = getRateLimiter(RateLimiterMode.Preview, token);
break;
default:
rateLimiter = getRateLimiter(RateLimiterMode.Crawl, token);
break;
// case RateLimiterMode.Search:
// rateLimiter = await searchRateLimiter(RateLimiterMode.Search, token);
// break;
}
} }
const team_endpoint_token = const team_endpoint_token =
@ -307,8 +416,8 @@ export async function supaAuthenticateUser(
logger.error(`Rate limit exceeded: ${rateLimiterRes}`, { logger.error(`Rate limit exceeded: ${rateLimiterRes}`, {
teamId, teamId,
priceId, priceId,
plan: subscriptionData?.plan,
mode, mode,
rateLimits: chunk?.rate_limits,
rateLimiterRes, rateLimiterRes,
}); });
const secs = Math.round(rateLimiterRes.msBeforeNext / 1000) || 1; const secs = Math.round(rateLimiterRes.msBeforeNext / 1000) || 1;
@ -342,7 +451,6 @@ export async function supaAuthenticateUser(
success: true, success: true,
team_id: `preview_${iptoken}`, team_id: `preview_${iptoken}`,
chunk: null, chunk: null,
plan: "free",
}; };
// check the origin of the request and make sure its from firecrawl.dev // check the origin of the request and make sure its from firecrawl.dev
// const origin = req.headers.origin; // const origin = req.headers.origin;
@ -356,65 +464,9 @@ export async function supaAuthenticateUser(
// return { success: false, error: "Unauthorized: Invalid token", status: 401 }; // return { success: false, error: "Unauthorized: Invalid token", status: 401 };
} }
if (token && isTestSuiteToken(token)) {
return {
success: true,
team_id: teamId ?? undefined,
// Now we have a test suite plan
plan: "testSuite",
chunk,
};
}
return { return {
success: true, success: true,
team_id: teamId ?? undefined, team_id: teamId ?? undefined,
plan: (subscriptionData?.plan ?? "") as PlanType,
chunk, chunk,
}; };
} }
function getPlanByPriceId(price_id: string | null): PlanType {
switch (price_id) {
case process.env.STRIPE_PRICE_ID_STARTER:
return "starter";
case process.env.STRIPE_PRICE_ID_STANDARD:
return "standard";
case process.env.STRIPE_PRICE_ID_SCALE:
return "scale";
case process.env.STRIPE_PRICE_ID_HOBBY:
case process.env.STRIPE_PRICE_ID_HOBBY_YEARLY:
return "hobby";
case process.env.STRIPE_PRICE_ID_STANDARD_NEW:
case process.env.STRIPE_PRICE_ID_STANDARD_NEW_YEARLY:
return "standardnew";
case process.env.STRIPE_PRICE_ID_GROWTH:
case process.env.STRIPE_PRICE_ID_GROWTH_YEARLY:
case process.env.STRIPE_PRICE_ID_SCALE_2M:
return "growth";
case process.env.STRIPE_PRICE_ID_GROWTH_DOUBLE_MONTHLY:
return "growthdouble";
case process.env.STRIPE_PRICE_ID_ETIER2C:
return "etier2c";
case process.env.STRIPE_PRICE_ID_ETIER1A_MONTHLY: //ocqh
return "etier1a";
case process.env.STRIPE_PRICE_ID_ETIER_SCALE_1_MONTHLY:
case process.env.STRIPE_PRICE_ID_ETIER_SCALE_1_YEARLY:
case process.env.STRIPE_PRICE_ID_ETIER_SCALE_1_YEARLY_FIRECRAWL:
return "etierscale1";
case process.env.STRIPE_PRICE_ID_ETIER_SCALE_2_YEARLY:
case process.env.STRIPE_PRICE_ID_ETIER_SCALE_2_MONTHLY:
return "etierscale2";
case process.env.STRIPE_PRICE_ID_EXTRACT_STARTER_MONTHLY:
case process.env.STRIPE_PRICE_ID_EXTRACT_STARTER_YEARLY:
return "extract_starter";
case process.env.STRIPE_PRICE_ID_EXTRACT_EXPLORER_MONTHLY:
case process.env.STRIPE_PRICE_ID_EXTRACT_EXPLORER_YEARLY:
return "extract_explorer";
case process.env.STRIPE_PRICE_ID_EXTRACT_PRO_MONTHLY:
case process.env.STRIPE_PRICE_ID_EXTRACT_PRO_YEARLY:
return "extract_pro";
default:
return "free";
}
}

3
apps/api/src/controllers/v0/admin/acuc-cache-clear.ts
View File

@ -1,6 +1,6 @@
import { Request, Response } from "express"; import { Request, Response } from "express";
import { supabase_service } from "../../../services/supabase"; import { supabase_service } from "../../../services/supabase";
import { clearACUC } from "../../auth"; import { clearACUC, clearACUCTeam } from "../../auth";
import { logger } from "../../../lib/logger"; import { logger } from "../../../lib/logger";
export async function acucCacheClearController(req: Request, res: Response) { export async function acucCacheClearController(req: Request, res: Response) {
@ -13,6 +13,7 @@ export async function acucCacheClearController(req: Request, res: Response) {
.eq("team_id", team_id); .eq("team_id", team_id);
await Promise.all((keys.data ?? []).map((x) => clearACUC(x.key))); await Promise.all((keys.data ?? []).map((x) => clearACUC(x.key)));
await clearACUCTeam(team_id);
logger.info(`ACUC cache cleared for team ${team_id}`); logger.info(`ACUC cache cleared for team ${team_id}`);
res.json({ ok: true }); res.json({ ok: true });

8
apps/api/src/controllers/v0/crawl.ts
View File

@ -39,7 +39,7 @@ export async function crawlController(req: Request, res: Response) {
return res.status(auth.status).json({ error: auth.error }); return res.status(auth.status).json({ error: auth.error });
} }
const { team_id, plan, chunk } = auth; const { team_id, chunk } = auth;
redisConnection.sadd("teams_using_v0", team_id) redisConnection.sadd("teams_using_v0", team_id)
.catch(error => logger.error("Failed to add team to teams_using_v0", { error, team_id })); .catch(error => logger.error("Failed to add team to teams_using_v0", { error, team_id }));
@ -170,7 +170,6 @@ export async function crawlController(req: Request, res: Response) {
scrapeOptions, scrapeOptions,
internalOptions, internalOptions,
team_id, team_id,
plan,
createdAt: Date.now(), createdAt: Date.now(),
}; };
@ -190,7 +189,6 @@ export async function crawlController(req: Request, res: Response) {
if (urls.length === 0) return; if (urls.length === 0) return;
let jobPriority = await getJobPriority({ let jobPriority = await getJobPriority({
plan,
team_id, team_id,
basePriority: 21, basePriority: 21,
}); });
@ -205,7 +203,6 @@ export async function crawlController(req: Request, res: Response) {
scrapeOptions, scrapeOptions,
internalOptions, internalOptions,
team_id, team_id,
plan,
origin: req.body.origin ?? defaultOrigin, origin: req.body.origin ?? defaultOrigin,
crawl_id: id, crawl_id: id,
sitemapped: true, sitemapped: true,
@ -236,7 +233,7 @@ export async function crawlController(req: Request, res: Response) {
await lockURL(id, sc, url); await lockURL(id, sc, url);
// Not needed, first one should be 15. // Not needed, first one should be 15.
// const jobPriority = await getJobPriority({plan, team_id, basePriority: 10}) // const jobPriority = await getJobPriority({team_id, basePriority: 10})
const jobId = uuidv4(); const jobId = uuidv4();
await addScrapeJob( await addScrapeJob(
@ -247,7 +244,6 @@ export async function crawlController(req: Request, res: Response) {
scrapeOptions, scrapeOptions,
internalOptions, internalOptions,
team_id, team_id,
plan: plan!,
origin: req.body.origin ?? defaultOrigin, origin: req.body.origin ?? defaultOrigin,
crawl_id: id, crawl_id: id,
}, },

5
apps/api/src/controllers/v0/crawlPreview.ts
View File

@ -31,8 +31,6 @@ export async function crawlPreviewController(req: Request, res: Response) {
return res.status(auth.status).json({ error: auth.error }); return res.status(auth.status).json({ error: auth.error });
} }
const { plan } = auth;
let url = req.body.url; let url = req.body.url;
if (!url) { if (!url) {
return res.status(400).json({ error: "Url is required" }); return res.status(400).json({ error: "Url is required" });
@ -108,7 +106,6 @@ export async function crawlPreviewController(req: Request, res: Response) {
scrapeOptions, scrapeOptions,
internalOptions, internalOptions,
team_id, team_id,
plan,
robots, robots,
createdAt: Date.now(), createdAt: Date.now(),
}; };
@ -130,7 +127,6 @@ export async function crawlPreviewController(req: Request, res: Response) {
url, url,
mode: "single_urls", mode: "single_urls",
team_id, team_id,
plan: plan!,
crawlerOptions, crawlerOptions,
scrapeOptions, scrapeOptions,
internalOptions, internalOptions,
@ -153,7 +149,6 @@ export async function crawlPreviewController(req: Request, res: Response) {
url, url,
mode: "single_urls", mode: "single_urls",
team_id, team_id,
plan: plan!,
crawlerOptions, crawlerOptions,
scrapeOptions, scrapeOptions,
internalOptions, internalOptions,

9
apps/api/src/controllers/v0/scrape.ts
View File

@ -5,7 +5,7 @@ import {
checkTeamCredits, checkTeamCredits,
} from "../../services/billing/credit_billing"; } from "../../services/billing/credit_billing";
import { authenticateUser } from "../auth"; import { authenticateUser } from "../auth";
import { PlanType, RateLimiterMode } from "../../types"; import { RateLimiterMode } from "../../types";
import { logJob } from "../../services/logging/log_job"; import { logJob } from "../../services/logging/log_job";
import { import {
fromLegacyCombo, fromLegacyCombo,
@ -39,7 +39,6 @@ export async function scrapeHelper(
pageOptions: PageOptions, pageOptions: PageOptions,
extractorOptions: ExtractorOptions, extractorOptions: ExtractorOptions,
timeout: number, timeout: number,
plan?: PlanType,
): Promise<{ ): Promise<{
success: boolean; success: boolean;
error?: string; error?: string;
@ -59,7 +58,7 @@ export async function scrapeHelper(
}; };
} }
const jobPriority = await getJobPriority({ plan, team_id, basePriority: 10 }); const jobPriority = await getJobPriority({ team_id, basePriority: 10 });
const { scrapeOptions, internalOptions } = fromLegacyCombo( const { scrapeOptions, internalOptions } = fromLegacyCombo(
pageOptions, pageOptions,
@ -76,7 +75,6 @@ export async function scrapeHelper(
team_id, team_id,
scrapeOptions, scrapeOptions,
internalOptions, internalOptions,
plan: plan!,
origin: req.body.origin ?? defaultOrigin, origin: req.body.origin ?? defaultOrigin,
is_scrape: true, is_scrape: true,
}, },
@ -180,7 +178,7 @@ export async function scrapeController(req: Request, res: Response) {
return res.status(auth.status).json({ error: auth.error }); return res.status(auth.status).json({ error: auth.error });
} }
const { team_id, plan, chunk } = auth; const { team_id, chunk } = auth;
redisConnection.sadd("teams_using_v0", team_id) redisConnection.sadd("teams_using_v0", team_id)
.catch(error => logger.error("Failed to add team to teams_using_v0", { error, team_id })); .catch(error => logger.error("Failed to add team to teams_using_v0", { error, team_id }));
@ -240,7 +238,6 @@ export async function scrapeController(req: Request, res: Response) {
pageOptions, pageOptions,
extractorOptions, extractorOptions,
timeout, timeout,
plan,
); );
const endTime = new Date().getTime(); const endTime = new Date().getTime();
const timeTakenInSeconds = (endTime - startTime) / 1000; const timeTakenInSeconds = (endTime - startTime) / 1000;

8
apps/api/src/controllers/v0/search.ts
View File

@ -4,7 +4,7 @@ import {
checkTeamCredits, checkTeamCredits,
} from "../../services/billing/credit_billing"; } from "../../services/billing/credit_billing";
import { authenticateUser } from "../auth"; import { authenticateUser } from "../auth";
import { PlanType, RateLimiterMode } from "../../types"; import { RateLimiterMode } from "../../types";
import { logJob } from "../../services/logging/log_job"; import { logJob } from "../../services/logging/log_job";
import { PageOptions, SearchOptions } from "../../lib/entities"; import { PageOptions, SearchOptions } from "../../lib/entities";
import { search } from "../../search"; import { search } from "../../search";
@ -31,7 +31,6 @@ export async function searchHelper(
crawlerOptions: any, crawlerOptions: any,
pageOptions: PageOptions, pageOptions: PageOptions,
searchOptions: SearchOptions, searchOptions: SearchOptions,
plan: PlanType | undefined,
): Promise<{ ): Promise<{
success: boolean; success: boolean;
error?: string; error?: string;
@ -94,7 +93,7 @@ export async function searchHelper(
return { success: true, error: "No search results found", returnCode: 200 }; return { success: true, error: "No search results found", returnCode: 200 };
} }
const jobPriority = await getJobPriority({ plan, team_id, basePriority: 20 }); const jobPriority = await getJobPriority({ team_id, basePriority: 20 });
// filter out social media links // filter out social media links
@ -163,7 +162,7 @@ export async function searchController(req: Request, res: Response) {
if (!auth.success) { if (!auth.success) {
return res.status(auth.status).json({ error: auth.error }); return res.status(auth.status).json({ error: auth.error });
} }
const { team_id, plan, chunk } = auth; const { team_id, chunk } = auth;
redisConnection.sadd("teams_using_v0", team_id) redisConnection.sadd("teams_using_v0", team_id)
.catch(error => logger.error("Failed to add team to teams_using_v0", { error, team_id })); .catch(error => logger.error("Failed to add team to teams_using_v0", { error, team_id }));
@ -202,7 +201,6 @@ export async function searchController(req: Request, res: Response) {
crawlerOptions, crawlerOptions,
pageOptions, pageOptions,
searchOptions, searchOptions,
plan,
); );
const endTime = new Date().getTime(); const endTime = new Date().getTime();
const timeTakenInSeconds = (endTime - startTime) / 1000; const timeTakenInSeconds = (endTime - startTime) / 1000;

4
apps/api/src/controllers/v1/batch-scrape.ts
View File

@ -40,7 +40,6 @@ export async function batchScrapeController(
module: "api/v1", module: "api/v1",
method: "batchScrapeController", method: "batchScrapeController",
teamId: req.auth.team_id, teamId: req.auth.team_id,
plan: req.auth.plan,
}); });
let urls = req.body.urls; let urls = req.body.urls;
@ -85,7 +84,6 @@ export async function batchScrapeController(
internalOptions: { disableSmartWaitCache: true, teamId: req.auth.team_id }, // NOTE: smart wait disabled for batch scrapes to ensure contentful scrape, speed does not matter internalOptions: { disableSmartWaitCache: true, teamId: req.auth.team_id }, // NOTE: smart wait disabled for batch scrapes to ensure contentful scrape, speed does not matter
team_id: req.auth.team_id, team_id: req.auth.team_id,
createdAt: Date.now(), createdAt: Date.now(),
plan: req.auth.plan,
}; };
if (!req.body.appendToId) { if (!req.body.appendToId) {
@ -99,7 +97,6 @@ export async function batchScrapeController(
if (urls.length > 1000) { if (urls.length > 1000) {
// set base to 21 // set base to 21
jobPriority = await getJobPriority({ jobPriority = await getJobPriority({
plan: req.auth.plan,
team_id: req.auth.team_id, team_id: req.auth.team_id,
basePriority: 21, basePriority: 21,
}); });
@ -116,7 +113,6 @@ export async function batchScrapeController(
url: x, url: x,
mode: "single_urls" as const, mode: "single_urls" as const,
team_id: req.auth.team_id, team_id: req.auth.team_id,
plan: req.auth.plan!,
crawlerOptions: null, crawlerOptions: null,
scrapeOptions, scrapeOptions,
origin: "api", origin: "api",

10
apps/api/src/controllers/v1/concurrency-check.ts
View File

@ -1,13 +1,10 @@
import { authenticateUser } from "../auth";
import { import {
ConcurrencyCheckParams, ConcurrencyCheckParams,
ConcurrencyCheckResponse, ConcurrencyCheckResponse,
RequestWithAuth, RequestWithAuth,
} from "./types"; } from "./types";
import { RateLimiterMode, PlanType } from "../../types";
import { Response } from "express"; import { Response } from "express";
import { redisConnection } from "../../services/queue-service"; import { redisConnection } from "../../services/queue-service";
import { getConcurrencyLimitMax } from "../../services/rate-limiter";
// Basically just middleware and error wrapping // Basically just middleware and error wrapping
export async function concurrencyCheckController( export async function concurrencyCheckController(
@ -22,14 +19,9 @@ export async function concurrencyCheckController(
Infinity, Infinity,
); );
const maxConcurrency = getConcurrencyLimitMax(
req.auth.plan as PlanType,
req.auth.team_id,
);
return res.status(200).json({ return res.status(200).json({
success: true, success: true,
concurrency: activeJobsOfTeam.length, concurrency: activeJobsOfTeam.length,
maxConcurrency: maxConcurrency, maxConcurrency: req.acuc.concurrency,
}); });
} }

4
apps/api/src/controllers/v1/crawl-status-ws.ts
View File

@ -191,9 +191,9 @@ export async function crawlStatusWSController(
}); });
} }
const { team_id, plan } = auth; const { team_id } = auth;
req.auth = { team_id, plan }; req.auth = { team_id };
await crawlStatusWS(ws, req); await crawlStatusWS(ws, req);
} catch (err) { } catch (err) {

3
apps/api/src/controllers/v1/crawl.ts
View File

@ -25,7 +25,6 @@ export async function crawlController(
module: "api/v1", module: "api/v1",
method: "crawlController", method: "crawlController",
teamId: req.auth.team_id, teamId: req.auth.team_id,
plan: req.auth.plan,
}); });
logger.debug("Crawl " + id + " starting", { logger.debug("Crawl " + id + " starting", {
request: req.body, request: req.body,
@ -84,7 +83,6 @@ export async function crawlController(
internalOptions: { disableSmartWaitCache: true, teamId: req.auth.team_id }, // NOTE: smart wait disabled for crawls to ensure contentful scrape, speed does not matter internalOptions: { disableSmartWaitCache: true, teamId: req.auth.team_id }, // NOTE: smart wait disabled for crawls to ensure contentful scrape, speed does not matter
team_id: req.auth.team_id, team_id: req.auth.team_id,
createdAt: Date.now(), createdAt: Date.now(),
plan: req.auth.plan,
}; };
const crawler = crawlToCrawler(id, sc); const crawler = crawlToCrawler(id, sc);
@ -104,7 +102,6 @@ export async function crawlController(
url: req.body.url, url: req.body.url,
mode: "kickoff" as const, mode: "kickoff" as const,
team_id: req.auth.team_id, team_id: req.auth.team_id,
plan: req.auth.plan,
crawlerOptions, crawlerOptions,
scrapeOptions: sc.scrapeOptions, scrapeOptions: sc.scrapeOptions,
internalOptions: sc.internalOptions, internalOptions: sc.internalOptions,

4
apps/api/src/controllers/v1/credit-usage.ts
View File

@ -1,6 +1,6 @@
import { Request, Response } from "express"; import { Request, Response } from "express";
import { RequestWithAuth } from "./types"; import { RequestWithAuth } from "./types";
import { getACUC } from "../auth"; import { getACUCTeam } from "../auth";
import { logger } from "../../lib/logger"; import { logger } from "../../lib/logger";
export async function creditUsageController( export async function creditUsageController(
@ -20,7 +20,7 @@ export async function creditUsageController(
} }
// Otherwise fetch fresh data // Otherwise fetch fresh data
const chunk = await getACUC(req.auth.team_id); const chunk = await getACUCTeam(req.auth.team_id);
if (!chunk) { if (!chunk) {
res.status(404).json({ res.status(404).json({
success: false, success: false,

2
apps/api/src/controllers/v1/deep-research.ts
View File

@ -52,7 +52,6 @@ export async function deepResearchController(
const jobData = { const jobData = {
request: req.body, request: req.body,
teamId: req.auth.team_id, teamId: req.auth.team_id,
plan: req.auth.plan,
subId: req.acuc?.sub_id, subId: req.acuc?.sub_id,
researchId, researchId,
}; };
@ -60,7 +59,6 @@ export async function deepResearchController(
await saveDeepResearch(researchId, { await saveDeepResearch(researchId, {
id: researchId, id: researchId,
team_id: req.auth.team_id, team_id: req.auth.team_id,
plan: req.auth.plan,
createdAt: Date.now(), createdAt: Date.now(),
status: "processing", status: "processing",
currentDepth: 0, currentDepth: 0,

3
apps/api/src/controllers/v1/extract.ts
View File

@ -22,7 +22,6 @@ export async function oldExtract(
const result = await performExtraction(extractId, { const result = await performExtraction(extractId, {
request: req.body, request: req.body,
teamId: req.auth.team_id, teamId: req.auth.team_id,
plan: req.auth.plan ?? "free",
subId: req.acuc?.sub_id ?? undefined, subId: req.acuc?.sub_id ?? undefined,
}); });
@ -52,7 +51,6 @@ export async function extractController(
const jobData = { const jobData = {
request: req.body, request: req.body,
teamId: req.auth.team_id, teamId: req.auth.team_id,
plan: req.auth.plan,
subId: req.acuc?.sub_id, subId: req.acuc?.sub_id,
extractId, extractId,
}; };
@ -68,7 +66,6 @@ export async function extractController(
await saveExtract(extractId, { await saveExtract(extractId, {
id: extractId, id: extractId,
team_id: req.auth.team_id, team_id: req.auth.team_id,
plan: req.auth.plan,
createdAt: Date.now(), createdAt: Date.now(),
status: "processing", status: "processing",
showSteps: req.body.__experimental_streamSteps, showSteps: req.body.__experimental_streamSteps,

2
apps/api/src/controllers/v1/generate-llmstxt.ts
View File

@ -30,7 +30,6 @@ export async function generateLLMsTextController(
const jobData = { const jobData = {
request: req.body, request: req.body,
teamId: req.auth.team_id, teamId: req.auth.team_id,
plan: req.auth.plan,
subId: req.acuc?.sub_id, subId: req.acuc?.sub_id,
generationId, generationId,
}; };
@ -38,7 +37,6 @@ export async function generateLLMsTextController(
await saveGeneratedLlmsTxt(generationId, { await saveGeneratedLlmsTxt(generationId, {
id: generationId, id: generationId,
team_id: req.auth.team_id, team_id: req.auth.team_id,
plan: req.auth.plan!, // Add non-null assertion since plan is required
createdAt: Date.now(), createdAt: Date.now(),
status: "processing", status: "processing",
url: req.body.url, url: req.body.url,

4
apps/api/src/controllers/v1/map.ts
View File

@ -50,7 +50,6 @@ export async function getMapResults({
includeSubdomains = true, includeSubdomains = true,
crawlerOptions = {}, crawlerOptions = {},
teamId, teamId,
plan,
origin, origin,
includeMetadata = false, includeMetadata = false,
allowExternalLinks, allowExternalLinks,
@ -65,7 +64,6 @@ export async function getMapResults({
includeSubdomains?: boolean; includeSubdomains?: boolean;
crawlerOptions?: any; crawlerOptions?: any;
teamId: string; teamId: string;
plan?: string;
origin?: string; origin?: string;
includeMetadata?: boolean; includeMetadata?: boolean;
allowExternalLinks?: boolean; allowExternalLinks?: boolean;
@ -88,7 +86,6 @@ export async function getMapResults({
internalOptions: { teamId }, internalOptions: { teamId },
team_id: teamId, team_id: teamId,
createdAt: Date.now(), createdAt: Date.now(),
plan: plan,
}; };
const crawler = crawlToCrawler(id, sc); const crawler = crawlToCrawler(id, sc);
@ -322,7 +319,6 @@ export async function mapController(
crawlerOptions: req.body, crawlerOptions: req.body,
origin: req.body.origin, origin: req.body.origin,
teamId: req.auth.team_id, teamId: req.auth.team_id,
plan: req.auth.plan,
abort: abort.signal, abort: abort.signal,
mock: req.body.useMock, mock: req.body.useMock,
filterByPath: req.body.filterByPath !== false, filterByPath: req.body.filterByPath !== false,

1
apps/api/src/controllers/v1/scrape-status.ts
View File

@ -1,6 +1,5 @@
import { Response } from "express"; import { Response } from "express";
import { supabaseGetJobByIdOnlyData } from "../../lib/supabase-jobs"; import { supabaseGetJobByIdOnlyData } from "../../lib/supabase-jobs";
import { scrapeStatusRateLimiter } from "../../services/rate-limiter";
export async function scrapeStatusController(req: any, res: any) { export async function scrapeStatusController(req: any, res: any) {
const allowedTeams = [ const allowedTeams = [

3
apps/api/src/controllers/v1/scrape.ts
View File

@ -12,7 +12,6 @@ import { v4 as uuidv4 } from "uuid";
import { addScrapeJob, waitForJob } from "../../services/queue-jobs"; import { addScrapeJob, waitForJob } from "../../services/queue-jobs";
import { logJob } from "../../services/logging/log_job"; import { logJob } from "../../services/logging/log_job";
import { getJobPriority } from "../../lib/job-priority"; import { getJobPriority } from "../../lib/job-priority";
import { PlanType } from "../../types";
import { getScrapeQueue } from "../../services/queue-service"; import { getScrapeQueue } from "../../services/queue-service";
export async function scrapeController( export async function scrapeController(
@ -38,7 +37,6 @@ export async function scrapeController(
const startTime = new Date().getTime(); const startTime = new Date().getTime();
const jobPriority = await getJobPriority({ const jobPriority = await getJobPriority({
plan: req.auth.plan as PlanType,
team_id: req.auth.team_id, team_id: req.auth.team_id,
basePriority: 10, basePriority: 10,
}); });
@ -51,7 +49,6 @@ export async function scrapeController(
team_id: req.auth.team_id, team_id: req.auth.team_id,
scrapeOptions: req.body, scrapeOptions: req.body,
internalOptions: { teamId: req.auth.team_id }, internalOptions: { teamId: req.auth.team_id },
plan: req.auth.plan!,
origin: req.body.origin, origin: req.body.origin,
is_scrape: true, is_scrape: true,
}, },

7
apps/api/src/controllers/v1/search.ts
View File

@ -13,7 +13,7 @@ import { v4 as uuidv4 } from "uuid";
import { addScrapeJob, waitForJob } from "../../services/queue-jobs"; import { addScrapeJob, waitForJob } from "../../services/queue-jobs";
import { logJob } from "../../services/logging/log_job"; import { logJob } from "../../services/logging/log_job";
import { getJobPriority } from "../../lib/job-priority"; import { getJobPriority } from "../../lib/job-priority";
import { PlanType, Mode } from "../../types"; import { Mode } from "../../types";
import { getScrapeQueue } from "../../services/queue-service"; import { getScrapeQueue } from "../../services/queue-service";
import { search } from "../../search"; import { search } from "../../search";
import { isUrlBlocked } from "../../scraper/WebScraper/utils/blocklist"; import { isUrlBlocked } from "../../scraper/WebScraper/utils/blocklist";
@ -25,7 +25,6 @@ export async function searchAndScrapeSearchResult(
query: string, query: string,
options: { options: {
teamId: string; teamId: string;
plan: PlanType | undefined;
origin: string; origin: string;
timeout: number; timeout: number;
scrapeOptions: ScrapeOptions; scrapeOptions: ScrapeOptions;
@ -60,7 +59,6 @@ async function scrapeSearchResult(
searchResult: { url: string; title: string; description: string }, searchResult: { url: string; title: string; description: string },
options: { options: {
teamId: string; teamId: string;
plan: PlanType | undefined;
origin: string; origin: string;
timeout: number; timeout: number;
scrapeOptions: ScrapeOptions; scrapeOptions: ScrapeOptions;
@ -68,7 +66,6 @@ async function scrapeSearchResult(
): Promise<Document> { ): Promise<Document> {
const jobId = uuidv4(); const jobId = uuidv4();
const jobPriority = await getJobPriority({ const jobPriority = await getJobPriority({
plan: options.plan as PlanType,
team_id: options.teamId, team_id: options.teamId,
basePriority: 10, basePriority: 10,
}); });
@ -84,7 +81,6 @@ async function scrapeSearchResult(
team_id: options.teamId, team_id: options.teamId,
scrapeOptions: options.scrapeOptions, scrapeOptions: options.scrapeOptions,
internalOptions: { teamId: options.teamId }, internalOptions: { teamId: options.teamId },
plan: options.plan || "free",
origin: options.origin, origin: options.origin,
is_scrape: true, is_scrape: true,
}, },
@ -190,7 +186,6 @@ export async function searchController(
const scrapePromises = searchResults.map((result) => const scrapePromises = searchResults.map((result) =>
scrapeSearchResult(result, { scrapeSearchResult(result, {
teamId: req.auth.team_id, teamId: req.auth.team_id,
plan: req.auth.plan,
origin: req.body.origin, origin: req.body.origin,
timeout: req.body.timeout, timeout: req.body.timeout,
scrapeOptions: req.body.scrapeOptions, scrapeOptions: req.body.scrapeOptions,

4
apps/api/src/controllers/v1/token-usage.ts
View File

@ -1,6 +1,6 @@
import { Request, Response } from "express"; import { Request, Response } from "express";
import { RequestWithAuth } from "./types"; import { RequestWithAuth } from "./types";
import { getACUC } from "../auth"; import { getACUC, getACUCTeam } from "../auth";
import { logger } from "../../lib/logger"; import { logger } from "../../lib/logger";
import { RateLimiterMode } from "../../types"; import { RateLimiterMode } from "../../types";
@ -21,7 +21,7 @@ export async function tokenUsageController(
} }
// Otherwise fetch fresh data // Otherwise fetch fresh data
const chunk = await getACUC(req.auth.team_id, false, true, RateLimiterMode.Extract); const chunk = await getACUCTeam(req.auth.team_id, false, true, RateLimiterMode.Extract);
if (!chunk) { if (!chunk) {
res.status(404).json({ res.status(404).json({
success: false, success: false,

24
apps/api/src/controllers/v1/types.ts
View File

@ -2,7 +2,6 @@ import { Request, Response } from "express";
import { z } from "zod"; import { z } from "zod";
import { isUrlBlocked } from "../../scraper/WebScraper/utils/blocklist"; import { isUrlBlocked } from "../../scraper/WebScraper/utils/blocklist";
import { protocolIncluded, checkUrl } from "../../lib/validateUrl"; import { protocolIncluded, checkUrl } from "../../lib/validateUrl";
import { PlanType } from "../../types";
import { countries } from "../../lib/validate-country"; import { countries } from "../../lib/validate-country";
import { import {
ExtractorOptions, ExtractorOptions,
@ -729,7 +728,6 @@ export type CrawlErrorsResponse =
type AuthObject = { type AuthObject = {
team_id: string; team_id: string;
plan: PlanType | undefined;
}; };
type Account = { type Account = {
@ -742,18 +740,36 @@ export type AuthCreditUsageChunk = {
sub_id: string | null; sub_id: string | null;
sub_current_period_start: string | null; sub_current_period_start: string | null;
sub_current_period_end: string | null; sub_current_period_end: string | null;
sub_user_id: string | null;
price_id: string | null; price_id: string | null;
price_credits: number; // credit limit with assoicated price, or free_credits (500) if free plan price_credits: number; // credit limit with assoicated price, or free_credits (500) if free plan
credits_used: number; credits_used: number;
coupon_credits: number; // do not rely on this number to be up to date after calling a billTeam coupon_credits: number; // do not rely on this number to be up to date after calling a billTeam
coupons: any[];
adjusted_credits_used: number; // credits this period minus coupons used adjusted_credits_used: number; // credits this period minus coupons used
remaining_credits: number; remaining_credits: number;
sub_user_id: string | null;
total_credits_sum: number; total_credits_sum: number;
plan_priority: {
bucketLimit: number;
planModifier: number;
};
rate_limits: {
crawl: number;
scrape: number;
search: number;
map: number;
extract: number;
preview: number;
crawlStatus: number;
extractStatus: number;
};
concurrency: number;
// appended on JS-side
is_extract?: boolean; is_extract?: boolean;
}; };
export type AuthCreditUsageChunkFromTeam = Omit<AuthCreditUsageChunk, "api_key">;
export interface RequestWithMaybeACUC< export interface RequestWithMaybeACUC<
ReqParams = {}, ReqParams = {},
ReqBody = undefined, ReqBody = undefined,

18
apps/api/src/lib/__tests__/job-priority.test.ts
View File

@ -4,7 +4,7 @@ import {
deleteJobPriority, deleteJobPriority,
} from "../job-priority"; } from "../job-priority";
import { redisConnection } from "../../services/queue-service"; import { redisConnection } from "../../services/queue-service";
import { PlanType } from "../../types"; import { } from "../../types";
jest.mock("../../services/queue-service", () => ({ jest.mock("../../services/queue-service", () => ({
redisConnection: { redisConnection: {
@ -46,14 +46,14 @@ describe("Job Priority Tests", () => {
test("getJobPriority should return correct priority based on plan and set length", async () => { test("getJobPriority should return correct priority based on plan and set length", async () => {
const team_id = "team1"; const team_id = "team1";
const plan: PlanType = "standard"; const plan = "standard";
(redisConnection.scard as jest.Mock).mockResolvedValue(150); (redisConnection.scard as jest.Mock).mockResolvedValue(150);
const priority = await getJobPriority({ plan, team_id }); const priority = await getJobPriority({ team_id });
expect(priority).toBe(10); expect(priority).toBe(10);
(redisConnection.scard as jest.Mock).mockResolvedValue(250); (redisConnection.scard as jest.Mock).mockResolvedValue(250);
const priorityExceeded = await getJobPriority({ plan, team_id }); const priorityExceeded = await getJobPriority({ team_id });
expect(priorityExceeded).toBe(20); // basePriority + Math.ceil((250 - 200) * 0.4) expect(priorityExceeded).toBe(20); // basePriority + Math.ceil((250 - 200) * 0.4)
}); });
@ -61,23 +61,23 @@ describe("Job Priority Tests", () => {
const team_id = "team1"; const team_id = "team1";
(redisConnection.scard as jest.Mock).mockResolvedValue(50); (redisConnection.scard as jest.Mock).mockResolvedValue(50);
let plan: PlanType = "hobby"; let plan = "hobby";
let priority = await getJobPriority({ plan, team_id }); let priority = await getJobPriority({ team_id });
expect(priority).toBe(10); expect(priority).toBe(10);
(redisConnection.scard as jest.Mock).mockResolvedValue(150); (redisConnection.scard as jest.Mock).mockResolvedValue(150);
plan = "hobby"; plan = "hobby";
priority = await getJobPriority({ plan, team_id }); priority = await getJobPriority({ team_id });
expect(priority).toBe(25); // basePriority + Math.ceil((150 - 50) * 0.3) expect(priority).toBe(25); // basePriority + Math.ceil((150 - 50) * 0.3)
(redisConnection.scard as jest.Mock).mockResolvedValue(25); (redisConnection.scard as jest.Mock).mockResolvedValue(25);
plan = "free"; plan = "free";
priority = await getJobPriority({ plan, team_id }); priority = await getJobPriority({ team_id });
expect(priority).toBe(10); expect(priority).toBe(10);
(redisConnection.scard as jest.Mock).mockResolvedValue(60); (redisConnection.scard as jest.Mock).mockResolvedValue(60);
plan = "free"; plan = "free";
priority = await getJobPriority({ plan, team_id }); priority = await getJobPriority({ team_id });
expect(priority).toBe(28); // basePriority + Math.ceil((60 - 25) * 0.5) expect(priority).toBe(28); // basePriority + Math.ceil((60 - 25) * 0.5)
}); });

4
apps/api/src/lib/crawl-redis.ts
View File

@ -11,7 +11,6 @@ export type StoredCrawl = {
scrapeOptions: Omit<ScrapeOptions, "timeout">; scrapeOptions: Omit<ScrapeOptions, "timeout">;
internalOptions: InternalOptions; internalOptions: InternalOptions;
team_id: string; team_id: string;
plan?: string;
robots?: string; robots?: string;
cancelled?: boolean; cancelled?: boolean;
createdAt: number; createdAt: number;
@ -24,7 +23,6 @@ export async function saveCrawl(id: string, crawl: StoredCrawl) {
method: "saveCrawl", method: "saveCrawl",
crawlId: id, crawlId: id,
teamId: crawl.team_id, teamId: crawl.team_id,
plan: crawl.plan,
}); });
await redisConnection.set("crawl:" + id, JSON.stringify(crawl)); await redisConnection.set("crawl:" + id, JSON.stringify(crawl));
await redisConnection.expire("crawl:" + id, 24 * 60 * 60); await redisConnection.expire("crawl:" + id, 24 * 60 * 60);
@ -274,7 +272,6 @@ export async function lockURL(
method: "lockURL", method: "lockURL",
preNormalizedURL: url, preNormalizedURL: url,
teamId: sc.team_id, teamId: sc.team_id,
plan: sc.plan,
}); });
if (typeof sc.crawlerOptions?.limit === "number") { if (typeof sc.crawlerOptions?.limit === "number") {
@ -335,7 +332,6 @@ export async function lockURLs(
module: "crawl-redis", module: "crawl-redis",
method: "lockURL", method: "lockURL",
teamId: sc.team_id, teamId: sc.team_id,
plan: sc.plan,
}); });
// Add to visited_unique set // Add to visited_unique set

1
apps/api/src/lib/deep-research/deep-research-redis.ts
View File

@ -32,7 +32,6 @@ export type DeepResearchFinding = {
export type StoredDeepResearch = { export type StoredDeepResearch = {
id: string; id: string;
team_id: string; team_id: string;
plan?: string;
createdAt: number; createdAt: number;
status: "processing" | "completed" | "failed" | "cancelled"; status: "processing" | "completed" | "failed" | "cancelled";
error?: any; error?: any;

6
apps/api/src/lib/deep-research/deep-research-service.ts
View File

@ -1,6 +1,5 @@
import { logger as _logger } from "../logger"; import { logger as _logger } from "../logger";
import { updateDeepResearch } from "./deep-research-redis"; import { updateDeepResearch } from "./deep-research-redis";
import { PlanType } from "../../types";
import { searchAndScrapeSearchResult } from "../../controllers/v1/search"; import { searchAndScrapeSearchResult } from "../../controllers/v1/search";
import { ResearchLLMService, ResearchStateManager } from "./research-manager"; import { ResearchLLMService, ResearchStateManager } from "./research-manager";
import { logJob } from "../../services/logging/log_job"; import { logJob } from "../../services/logging/log_job";
@ -10,7 +9,6 @@ import { ExtractOptions } from "../../controllers/v1/types";
interface DeepResearchServiceOptions { interface DeepResearchServiceOptions {
researchId: string; researchId: string;
teamId: string; teamId: string;
plan: string;
query: string; query: string;
maxDepth: number; maxDepth: number;
maxUrls: number; maxUrls: number;
@ -23,7 +21,7 @@ interface DeepResearchServiceOptions {
} }
export async function performDeepResearch(options: DeepResearchServiceOptions) { export async function performDeepResearch(options: DeepResearchServiceOptions) {
const { researchId, teamId, plan, timeLimit, subId, maxUrls } = options; const { researchId, teamId, timeLimit, subId, maxUrls } = options;
const startTime = Date.now(); const startTime = Date.now();
let currentTopic = options.query; let currentTopic = options.query;
let urlsAnalyzed = 0; let urlsAnalyzed = 0;
@ -39,7 +37,6 @@ export async function performDeepResearch(options: DeepResearchServiceOptions) {
const state = new ResearchStateManager( const state = new ResearchStateManager(
researchId, researchId,
teamId, teamId,
plan,
options.maxDepth, options.maxDepth,
logger, logger,
options.query, options.query,
@ -98,7 +95,6 @@ export async function performDeepResearch(options: DeepResearchServiceOptions) {
const response = await searchAndScrapeSearchResult(searchQuery.query, { const response = await searchAndScrapeSearchResult(searchQuery.query, {
teamId: options.teamId, teamId: options.teamId,
plan: options.plan as PlanType,
origin: "deep-research", origin: "deep-research",
timeout: 10000, timeout: 10000,
scrapeOptions: { scrapeOptions: {

1
apps/api/src/lib/deep-research/research-manager.ts
View File

@ -29,7 +29,6 @@ export class ResearchStateManager {
constructor( constructor(
private readonly researchId: string, private readonly researchId: string,
private readonly teamId: string, private readonly teamId: string,
private readonly plan: string,
private readonly maxDepth: number, private readonly maxDepth: number,
private readonly logger: Logger, private readonly logger: Logger,
private readonly topic: string, private readonly topic: string,

4
apps/api/src/lib/extract/document-scraper.ts
View File

@ -1,5 +1,4 @@
import { Document, ScrapeOptions, URLTrace, scrapeOptions } from "../../controllers/v1/types"; import { Document, ScrapeOptions, URLTrace, scrapeOptions } from "../../controllers/v1/types";
import { PlanType } from "../../types";
import { logger } from "../logger"; import { logger } from "../logger";
import { getScrapeQueue } from "../../services/queue-service"; import { getScrapeQueue } from "../../services/queue-service";
import { waitForJob } from "../../services/queue-jobs"; import { waitForJob } from "../../services/queue-jobs";
@ -10,7 +9,6 @@ import type { Logger } from "winston";
interface ScrapeDocumentOptions { interface ScrapeDocumentOptions {
url: string; url: string;
teamId: string; teamId: string;
plan: PlanType;
origin: string; origin: string;
timeout: number; timeout: number;
isSingleUrl?: boolean; isSingleUrl?: boolean;
@ -31,7 +29,6 @@ export async function scrapeDocument(
async function attemptScrape(timeout: number) { async function attemptScrape(timeout: number) {
const jobId = crypto.randomUUID(); const jobId = crypto.randomUUID();
const jobPriority = await getJobPriority({ const jobPriority = await getJobPriority({
plan: options.plan,
team_id: options.teamId, team_id: options.teamId,
basePriority: 10, basePriority: 10,
}); });
@ -46,7 +43,6 @@ export async function scrapeDocument(
useCache: true, useCache: true,
teamId: options.teamId, teamId: options.teamId,
}, },
plan: options.plan,
origin: options.origin, origin: options.origin,
is_scrape: true, is_scrape: true,
}, },

1
apps/api/src/lib/extract/extract-redis.ts
View File

@ -24,7 +24,6 @@ export type ExtractedStep = {
export type StoredExtract = { export type StoredExtract = {
id: string; id: string;
team_id: string; team_id: string;
plan?: string;
createdAt: number; createdAt: number;
status: "processing" | "completed" | "failed" | "cancelled"; status: "processing" | "completed" | "failed" | "cancelled";
error?: any; error?: any;

7
apps/api/src/lib/extract/extraction-service.ts
View File

@ -4,7 +4,6 @@ import {
TokenUsage, TokenUsage,
URLTrace, URLTrace,
} from "../../controllers/v1/types"; } from "../../controllers/v1/types";
import { PlanType } from "../../types";
import { logger as _logger } from "../logger"; import { logger as _logger } from "../logger";
import { generateBasicCompletion, processUrl } from "./url-processor"; import { generateBasicCompletion, processUrl } from "./url-processor";
import { scrapeDocument } from "./document-scraper"; import { scrapeDocument } from "./document-scraper";
@ -44,7 +43,6 @@ import { buildRephraseToSerpPrompt } from "./build-prompts";
interface ExtractServiceOptions { interface ExtractServiceOptions {
request: ExtractRequest; request: ExtractRequest;
teamId: string; teamId: string;
plan: PlanType;
subId?: string; subId?: string;
cacheMode?: "load" | "save" | "direct"; cacheMode?: "load" | "save" | "direct";
cacheKey?: string; cacheKey?: string;
@ -76,7 +74,7 @@ export async function performExtraction(
extractId: string, extractId: string,
options: ExtractServiceOptions, options: ExtractServiceOptions,
): Promise<ExtractResult> { ): Promise<ExtractResult> {
const { request, teamId, plan, subId } = options; const { request, teamId, subId } = options;
const urlTraces: URLTrace[] = []; const urlTraces: URLTrace[] = [];
let docsMap: Map<string, Document> = new Map(); let docsMap: Map<string, Document> = new Map();
let singleAnswerCompletions: completions | null = null; let singleAnswerCompletions: completions | null = null;
@ -161,7 +159,6 @@ export async function performExtraction(
url, url,
prompt: request.prompt, prompt: request.prompt,
teamId, teamId,
plan,
allowExternalLinks: request.allowExternalLinks, allowExternalLinks: request.allowExternalLinks,
origin: request.origin, origin: request.origin,
limit: request.limit, limit: request.limit,
@ -311,7 +308,6 @@ export async function performExtraction(
{ {
url, url,
teamId, teamId,
plan,
origin: request.origin || "api", origin: request.origin || "api",
timeout, timeout,
}, },
@ -574,7 +570,6 @@ export async function performExtraction(
{ {
url, url,
teamId, teamId,
plan,
origin: request.origin || "api", origin: request.origin || "api",
timeout, timeout,
}, },

4
apps/api/src/lib/extract/url-processor.ts
View File

@ -1,6 +1,5 @@
import { MapDocument, URLTrace } from "../../controllers/v1/types"; import { MapDocument, URLTrace } from "../../controllers/v1/types";
import { getMapResults } from "../../controllers/v1/map"; import { getMapResults } from "../../controllers/v1/map";
import { PlanType } from "../../types";
import { removeDuplicateUrls } from "../validateUrl"; import { removeDuplicateUrls } from "../validateUrl";
import { isUrlBlocked } from "../../scraper/WebScraper/utils/blocklist"; import { isUrlBlocked } from "../../scraper/WebScraper/utils/blocklist";
import { buildPreRerankPrompt, buildRefrasedPrompt } from "./build-prompts"; import { buildPreRerankPrompt, buildRefrasedPrompt } from "./build-prompts";
@ -23,7 +22,6 @@ interface ProcessUrlOptions {
prompt?: string; prompt?: string;
schema?: any; schema?: any;
teamId: string; teamId: string;
plan: PlanType;
allowExternalLinks?: boolean; allowExternalLinks?: boolean;
origin?: string; origin?: string;
limit?: number; limit?: number;
@ -80,7 +78,6 @@ export async function processUrl(
url: baseUrl, url: baseUrl,
search: searchQuery, search: searchQuery,
teamId: options.teamId, teamId: options.teamId,
plan: options.plan,
allowExternalLinks: options.allowExternalLinks, allowExternalLinks: options.allowExternalLinks,
origin: options.origin, origin: options.origin,
limit: options.limit, limit: options.limit,
@ -117,7 +114,6 @@ export async function processUrl(
const retryMapResults = await getMapResults({ const retryMapResults = await getMapResults({
url: baseUrl, url: baseUrl,
teamId: options.teamId, teamId: options.teamId,
plan: options.plan,
allowExternalLinks: options.allowExternalLinks, allowExternalLinks: options.allowExternalLinks,
origin: options.origin, origin: options.origin,
limit: options.limit, limit: options.limit,

1
apps/api/src/lib/generate-llmstxt/generate-llmstxt-redis.ts
View File

@ -4,7 +4,6 @@ import { logger as _logger } from "../logger";
export interface GenerationData { export interface GenerationData {
id: string; id: string;
team_id: string; team_id: string;
plan: string;
createdAt: number; createdAt: number;
status: "processing" | "completed" | "failed"; status: "processing" | "completed" | "failed";
url: string; url: string;

6
apps/api/src/lib/generate-llmstxt/generate-llmstxt-service.ts
View File

@ -3,7 +3,6 @@ import { updateGeneratedLlmsTxt } from "./generate-llmstxt-redis";
import { getMapResults } from "../../controllers/v1/map"; import { getMapResults } from "../../controllers/v1/map";
import { z } from "zod"; import { z } from "zod";
import { scrapeDocument } from "../extract/document-scraper"; import { scrapeDocument } from "../extract/document-scraper";
import { PlanType } from "../../types";
import { import {
getLlmsTextFromCache, getLlmsTextFromCache,
saveLlmsTextToCache, saveLlmsTextToCache,
@ -16,7 +15,6 @@ import { generateCompletions } from "../../scraper/scrapeURL/transformers/llmExt
interface GenerateLLMsTextServiceOptions { interface GenerateLLMsTextServiceOptions {
generationId: string; generationId: string;
teamId: string; teamId: string;
plan: PlanType;
url: string; url: string;
maxUrls: number; maxUrls: number;
showFullText: boolean; showFullText: boolean;
@ -64,7 +62,7 @@ function limitLlmsTxtEntries(llmstxt: string, maxEntries: number): string {
export async function performGenerateLlmsTxt( export async function performGenerateLlmsTxt(
options: GenerateLLMsTextServiceOptions, options: GenerateLLMsTextServiceOptions,
) { ) {
const { generationId, teamId, plan, url, maxUrls = 100, showFullText, subId } = const { generationId, teamId, url, maxUrls = 100, showFullText, subId } =
options; options;
const startTime = Date.now(); const startTime = Date.now();
const logger = _logger.child({ const logger = _logger.child({
@ -113,7 +111,6 @@ export async function performGenerateLlmsTxt(
const mapResult = await getMapResults({ const mapResult = await getMapResults({
url, url,
teamId, teamId,
plan,
limit: effectiveMaxUrls, limit: effectiveMaxUrls,
includeSubdomains: false, includeSubdomains: false,
ignoreSitemap: false, ignoreSitemap: false,
@ -142,7 +139,6 @@ export async function performGenerateLlmsTxt(
{ {
url, url,
teamId, teamId,
plan,
origin: url, origin: url,
timeout: 30000, timeout: 30000,
isSingleUrl: true, isSingleUrl: true,

52
apps/api/src/lib/job-priority.ts
View File

@ -1,5 +1,5 @@
import { redisConnection } from "../../src/services/queue-service"; import { getACUC, getACUCTeam } from "../controllers/auth";
import { PlanType } from "../../src/types"; import { redisConnection } from "../services/queue-service";
import { logger } from "./logger"; import { logger } from "./logger";
const SET_KEY_PREFIX = "limit_team_id:"; const SET_KEY_PREFIX = "limit_team_id:";
@ -29,11 +29,9 @@ export async function deleteJobPriority(team_id, job_id) {
} }
export async function getJobPriority({ export async function getJobPriority({
plan,
team_id, team_id,
basePriority = 10, basePriority = 10,
}: { }: {
plan: PlanType | undefined;
team_id: string; team_id: string;
basePriority?: number; basePriority?: number;
}): Promise<number> { }): Promise<number> {
@ -42,52 +40,16 @@ export async function getJobPriority({
} }
try { try {
const acuc = await getACUCTeam(team_id);
const setKey = SET_KEY_PREFIX + team_id; const setKey = SET_KEY_PREFIX + team_id;
// Get the length of the set // Get the length of the set
const setLength = await redisConnection.scard(setKey); const setLength = await redisConnection.scard(setKey);
// Determine the priority based on the plan and set length // Determine the priority based on the plan and set length
let planModifier = 1; let planModifier = acuc?.plan_priority.planModifier ?? 1;
let bucketLimit = 0; let bucketLimit = acuc?.plan_priority.bucketLimit ?? 25;
switch (plan) {
case "testSuite":
bucketLimit = 1000;
planModifier = 0.25;
break;
case "free":
bucketLimit = 25;
planModifier = 0.5;
break;
case "hobby":
bucketLimit = 100;
planModifier = 0.3;
break;
case "standard":
case "standardnew":
bucketLimit = 200;
planModifier = 0.2;
break;
case "growth":
case "growthdouble":
bucketLimit = 400;
planModifier = 0.1;
break;
case "etier2c":
bucketLimit = 1000;
planModifier = 0.05;
break;
case "etier1a":
bucketLimit = 1000;
planModifier = 0.05;
break;
default:
bucketLimit = 25;
planModifier = 1;
break;
}
// if length set is smaller than set, just return base priority // if length set is smaller than set, just return base priority
if (setLength <= bucketLimit) { if (setLength <= bucketLimit) {
@ -100,7 +62,7 @@ export async function getJobPriority({
} }
} catch (e) { } catch (e) {
logger.error( logger.error(
`Get job priority failed: ${team_id}, ${plan}, ${basePriority}`, `Get job priority failed: ${team_id}, ${basePriority}`,
); );
return basePriority; return basePriority;
} }

4
apps/api/src/routes/v1.ts
View File

@ -105,9 +105,9 @@ export function authMiddleware(
} }
} }
const { team_id, plan, chunk } = auth; const { team_id, chunk } = auth;
req.auth = { team_id, plan }; req.auth = { team_id };
req.acuc = chunk ?? undefined; req.acuc = chunk ?? undefined;
if (chunk) { if (chunk) {
req.account = { remainingCredits: chunk.remaining_credits }; req.account = { remainingCredits: chunk.remaining_credits };

14
apps/api/src/services/billing/batch_billing.ts
View File

@ -4,7 +4,7 @@ import { supabase_service } from "../supabase";
import * as Sentry from "@sentry/node"; import * as Sentry from "@sentry/node";
import { Queue } from "bullmq"; import { Queue } from "bullmq";
import { withAuth } from "../../lib/withAuth"; import { withAuth } from "../../lib/withAuth";
import { getACUC, setCachedACUC } from "../../controllers/auth"; import { getACUC, setCachedACUC, setCachedACUCTeam } from "../../controllers/auth";
// Configuration constants // Configuration constants
const BATCH_KEY = "billing_batch"; const BATCH_KEY = "billing_batch";
@ -298,7 +298,17 @@ async function supaBillTeam(
// Update cached ACUC to reflect the new credit usage // Update cached ACUC to reflect the new credit usage
(async () => { (async () => {
for (const apiKey of (data ?? []).map((x) => x.api_key)) { for (const apiKey of (data ?? []).map((x) => x.api_key)) {
await setCachedACUC(apiKey, (acuc) => await setCachedACUC(apiKey, is_extract, (acuc) =>
acuc
? {
...acuc,
credits_used: acuc.credits_used + credits,
adjusted_credits_used: acuc.adjusted_credits_used + credits,
remaining_credits: acuc.remaining_credits - credits,
}
: null,
);
await setCachedACUCTeam(team_id, is_extract, (acuc) =>
acuc acuc
? { ? {
...acuc, ...acuc,

13
apps/api/src/services/queue-jobs.ts
View File

@ -1,6 +1,6 @@
import { getScrapeQueue } from "./queue-service"; import { getScrapeQueue } from "./queue-service";
import { v4 as uuidv4 } from "uuid"; import { v4 as uuidv4 } from "uuid";
import { NotificationType, PlanType, WebScraperOptions } from "../types"; import { NotificationType, WebScraperOptions } from "../types";
import * as Sentry from "@sentry/node"; import * as Sentry from "@sentry/node";
import { import {
cleanOldConcurrencyLimitEntries, cleanOldConcurrencyLimitEntries,
@ -10,9 +10,9 @@ import {
pushConcurrencyLimitedJob, pushConcurrencyLimitedJob,
} from "../lib/concurrency-limit"; } from "../lib/concurrency-limit";
import { logger } from "../lib/logger"; import { logger } from "../lib/logger";
import { getConcurrencyLimitMax } from "./rate-limiter";
import { sendNotificationWithCustomDays } from './notification/email_notification'; import { sendNotificationWithCustomDays } from './notification/email_notification';
import { shouldSendConcurrencyLimitNotification } from './notification/notification-check'; import { shouldSendConcurrencyLimitNotification } from './notification/notification-check';
import { getACUC, getACUCTeam } from "../controllers/auth";
/** /**
* Checks if a job is a crawl or batch scrape based on its options * Checks if a job is a crawl or batch scrape based on its options
@ -51,8 +51,7 @@ export async function _addScrapeJobToBullMQ(
) { ) {
if ( if (
webScraperOptions && webScraperOptions &&
webScraperOptions.team_id && webScraperOptions.team_id
webScraperOptions.plan
) { ) {
await pushConcurrencyLimitActiveJob(webScraperOptions.team_id, jobId, 60 * 1000); // 60s default timeout await pushConcurrencyLimitActiveJob(webScraperOptions.team_id, jobId, 60 * 1000); // 60s default timeout
} }
@ -79,7 +78,7 @@ async function addScrapeJobRaw(
webScraperOptions.team_id webScraperOptions.team_id
) { ) {
const now = Date.now(); const now = Date.now();
maxConcurrency = getConcurrencyLimitMax(webScraperOptions.plan ?? "free", webScraperOptions.team_id); maxConcurrency = (await getACUCTeam(webScraperOptions.team_id))?.concurrency ?? 2;
cleanOldConcurrencyLimitEntries(webScraperOptions.team_id, now); cleanOldConcurrencyLimitEntries(webScraperOptions.team_id, now);
currentActiveConcurrency = (await getConcurrencyLimitActiveJobs(webScraperOptions.team_id, now)).length; currentActiveConcurrency = (await getConcurrencyLimitActiveJobs(webScraperOptions.team_id, now)).length;
concurrencyLimited = currentActiveConcurrency >= maxConcurrency; concurrencyLimited = currentActiveConcurrency >= maxConcurrency;
@ -170,9 +169,9 @@ export async function addScrapeJobs(
let currentActiveConcurrency = 0; let currentActiveConcurrency = 0;
let maxConcurrency = 0; let maxConcurrency = 0;
if (jobs[0].data && jobs[0].data.team_id && jobs[0].data.plan) { if (jobs[0].data && jobs[0].data.team_id) {
const now = Date.now(); const now = Date.now();
maxConcurrency = getConcurrencyLimitMax(jobs[0].data.plan as PlanType, jobs[0].data.team_id); maxConcurrency = (await getACUCTeam(jobs[0].data.team_id))?.concurrency ?? 2;
cleanOldConcurrencyLimitEntries(jobs[0].data.team_id, now); cleanOldConcurrencyLimitEntries(jobs[0].data.team_id, now);
currentActiveConcurrency = (await getConcurrencyLimitActiveJobs(jobs[0].data.team_id, now)).length; currentActiveConcurrency = (await getConcurrencyLimitActiveJobs(jobs[0].data.team_id, now)).length;

14
apps/api/src/services/queue-worker.ts
View File

@ -48,11 +48,9 @@ import {
deleteJobPriority, deleteJobPriority,
getJobPriority, getJobPriority,
} from "../../src/lib/job-priority"; } from "../../src/lib/job-priority";
import { PlanType, RateLimiterMode } from "../types";
import { getJobs } from "..//controllers/v1/crawl-status"; import { getJobs } from "..//controllers/v1/crawl-status";
import { configDotenv } from "dotenv"; import { configDotenv } from "dotenv";
import { scrapeOptions } from "../controllers/v1/types"; import { scrapeOptions } from "../controllers/v1/types";
import { getRateLimiterPoints } from "./rate-limiter";
import { import {
cleanOldConcurrencyLimitEntries, cleanOldConcurrencyLimitEntries,
pushConcurrencyLimitActiveJob, pushConcurrencyLimitActiveJob,
@ -144,7 +142,6 @@ async function finishCrawlIfNeeded(job: Job & { id: string }, sc: StoredCrawl) {
url, url,
mode: "single_urls" as const, mode: "single_urls" as const,
team_id: job.data.team_id, team_id: job.data.team_id,
plan: job.data.plan!,
crawlerOptions: { crawlerOptions: {
...job.data.crawlerOptions, ...job.data.crawlerOptions,
urlInvisibleInCurrentCrawl: true, urlInvisibleInCurrentCrawl: true,
@ -407,7 +404,6 @@ const processExtractJobInternal = async (
const result = await performExtraction(job.data.extractId, { const result = await performExtraction(job.data.extractId, {
request: job.data.request, request: job.data.request,
teamId: job.data.teamId, teamId: job.data.teamId,
plan: job.data.plan,
subId: job.data.subId, subId: job.data.subId,
}); });
@ -489,7 +485,6 @@ const processDeepResearchJobInternal = async (
const result = await performDeepResearch({ const result = await performDeepResearch({
researchId: job.data.researchId, researchId: job.data.researchId,
teamId: job.data.teamId, teamId: job.data.teamId,
plan: job.data.plan,
query: job.data.request.query, query: job.data.request.query,
maxDepth: job.data.request.maxDepth, maxDepth: job.data.request.maxDepth,
timeLimit: job.data.request.timeLimit, timeLimit: job.data.request.timeLimit,
@ -564,7 +559,6 @@ const processGenerateLlmsTxtJobInternal = async (
const result = await performGenerateLlmsTxt({ const result = await performGenerateLlmsTxt({
generationId: job.data.generationId, generationId: job.data.generationId,
teamId: job.data.teamId, teamId: job.data.teamId,
plan: job.data.plan,
url: job.data.request.url, url: job.data.request.url,
maxUrls: job.data.request.maxUrls, maxUrls: job.data.request.maxUrls,
showFullText: job.data.request.showFullText, showFullText: job.data.request.showFullText,
@ -682,7 +676,7 @@ const workerFun = async (
runningJobs.delete(job.id); runningJobs.delete(job.id);
} }
if (job.id && job.data && job.data.team_id && job.data.plan) { if (job.id && job.data && job.data.team_id) {
await removeConcurrencyLimitActiveJob(job.data.team_id, job.id); await removeConcurrencyLimitActiveJob(job.data.team_id, job.id);
cleanOldConcurrencyLimitEntries(job.data.team_id); cleanOldConcurrencyLimitEntries(job.data.team_id);
@ -805,7 +799,6 @@ async function processKickoffJob(job: Job & { id: string }, token: string) {
crawlerOptions: job.data.crawlerOptions, crawlerOptions: job.data.crawlerOptions,
scrapeOptions: scrapeOptions.parse(job.data.scrapeOptions), scrapeOptions: scrapeOptions.parse(job.data.scrapeOptions),
internalOptions: sc.internalOptions, internalOptions: sc.internalOptions,
plan: job.data.plan!,
origin: job.data.origin, origin: job.data.origin,
crawl_id: job.data.crawl_id, crawl_id: job.data.crawl_id,
webhook: job.data.webhook, webhook: job.data.webhook,
@ -844,7 +837,6 @@ async function processKickoffJob(job: Job & { id: string }, token: string) {
}); });
let jobPriority = await getJobPriority({ let jobPriority = await getJobPriority({
plan: job.data.plan,
team_id: job.data.team_id, team_id: job.data.team_id,
basePriority: 21, basePriority: 21,
}); });
@ -858,7 +850,6 @@ async function processKickoffJob(job: Job & { id: string }, token: string) {
url, url,
mode: "single_urls" as const, mode: "single_urls" as const,
team_id: job.data.team_id, team_id: job.data.team_id,
plan: job.data.plan!,
crawlerOptions: job.data.crawlerOptions, crawlerOptions: job.data.crawlerOptions,
scrapeOptions: job.data.scrapeOptions, scrapeOptions: job.data.scrapeOptions,
internalOptions: sc.internalOptions, internalOptions: sc.internalOptions,
@ -1155,7 +1146,6 @@ async function processJob(job: Job & { id: string }, token: string) {
if (await lockURL(job.data.crawl_id, sc, link)) { if (await lockURL(job.data.crawl_id, sc, link)) {
// This seems to work really welel // This seems to work really welel
const jobPriority = await getJobPriority({ const jobPriority = await getJobPriority({
plan: sc.plan as PlanType,
team_id: sc.team_id, team_id: sc.team_id,
basePriority: job.data.crawl_id ? 20 : 10, basePriority: job.data.crawl_id ? 20 : 10,
}); });
@ -1169,7 +1159,6 @@ async function processJob(job: Job & { id: string }, token: string) {
{ jobPriority, url: link }, { jobPriority, url: link },
); );
// console.log("plan: ", sc.plan);
// console.log("team_id: ", sc.team_id) // console.log("team_id: ", sc.team_id)
// console.log("base priority: ", job.data.crawl_id ? 20 : 10) // console.log("base priority: ", job.data.crawl_id ? 20 : 10)
// console.log("job priority: " , jobPriority, "\n\n\n") // console.log("job priority: " , jobPriority, "\n\n\n")
@ -1185,7 +1174,6 @@ async function processJob(job: Job & { id: string }, token: string) {
...sc.crawlerOptions, ...sc.crawlerOptions,
currentDiscoveryDepth: (job.data.crawlerOptions?.currentDiscoveryDepth ?? 0) + 1, currentDiscoveryDepth: (job.data.crawlerOptions?.currentDiscoveryDepth ?? 0) + 1,
}, },
plan: job.data.plan,
origin: job.data.origin, origin: job.data.origin,
crawl_id: job.data.crawl_id, crawl_id: job.data.crawl_id,
webhook: job.data.webhook, webhook: job.data.webhook,

642
apps/api/src/services/rate-limiter.test.ts
View File

@ -1,370 +1,370 @@
import { // import {
getRateLimiter, // getRateLimiter,
serverRateLimiter, // serverRateLimiter,
testSuiteRateLimiter, // redisRateLimitClient,
redisRateLimitClient, // } from "./rate-limiter";
} from "./rate-limiter"; // import { RateLimiterMode } from "../../src/types";
import { RateLimiterMode } from "../../src/types"; // import { RateLimiterRedis } from "rate-limiter-flexible";
import { RateLimiterRedis } from "rate-limiter-flexible";
describe("Rate Limiter Service", () => { // describe("Rate Limiter Service", () => {
beforeAll(async () => { // beforeAll(async () => {
try { // try {
await redisRateLimitClient.connect(); // await redisRateLimitClient.connect();
// if (process.env.REDIS_RATE_LIMIT_URL === "redis://localhost:6379") { // // if (process.env.REDIS_RATE_LIMIT_URL === "redis://localhost:6379") {
// console.log("Erasing all keys"); // // console.log("Erasing all keys");
// // erase all the keys that start with "test-prefix" // // // erase all the keys that start with "test-prefix"
// const keys = await redisRateLimitClient.keys("test-prefix:*"); // // const keys = await redisRateLimitClient.keys("test-prefix:*");
// if (keys.length > 0) { // // if (keys.length > 0) {
// await redisRateLimitClient.del(...keys); // // await redisRateLimitClient.del(...keys);
// } // // }
// } // // }
} catch (error) {} // } catch (error) {}
}); // });
afterAll(async () => { // afterAll(async () => {
try { // try {
// if (process.env.REDIS_RATE_LIMIT_URL === "redis://localhost:6379") { // // if (process.env.REDIS_RATE_LIMIT_URL === "redis://localhost:6379") {
await redisRateLimitClient.disconnect(); // await redisRateLimitClient.disconnect();
// } // // }
} catch (error) {} // } catch (error) {}
}); // });
it("should return the testSuiteRateLimiter for specific tokens", () => { // it("should return the testSuiteRateLimiter for specific tokens", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"crawl" as RateLimiterMode, // "crawl" as RateLimiterMode,
"test-prefix:a01ccae", // "test-prefix:a01ccae",
); // );
expect(limiter).toBe(testSuiteRateLimiter); // expect(limiter).toBe(testSuiteRateLimiter);
const limiter2 = getRateLimiter( // const limiter2 = getRateLimiter(
"scrape" as RateLimiterMode, // "scrape" as RateLimiterMode,
"test-prefix:6254cf9", // "test-prefix:6254cf9",
); // );
expect(limiter2).toBe(testSuiteRateLimiter); // expect(limiter2).toBe(testSuiteRateLimiter);
}); // });
it("should return the serverRateLimiter if mode is not found", () => { // it("should return the serverRateLimiter if mode is not found", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"nonexistent" as RateLimiterMode, // "nonexistent" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
); // );
expect(limiter.points).toBe(serverRateLimiter.points); // expect(limiter.points).toBe(serverRateLimiter.points);
}); // });
it("should return the correct rate limiter based on mode and plan", () => { // it("should return the correct rate limiter based on mode and plan", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"crawl" as RateLimiterMode, // "crawl" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"free", // "free",
); // );
expect(limiter.points).toBe(2); // expect(limiter.points).toBe(2);
const limiter2 = getRateLimiter( // const limiter2 = getRateLimiter(
"scrape" as RateLimiterMode, // "scrape" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"standard", // "standard",
); // );
expect(limiter2.points).toBe(100); // expect(limiter2.points).toBe(100);
const limiter3 = getRateLimiter( // const limiter3 = getRateLimiter(
"search" as RateLimiterMode, // "search" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"growth", // "growth",
); // );
expect(limiter3.points).toBe(500); // expect(limiter3.points).toBe(500);
const limiter4 = getRateLimiter( // const limiter4 = getRateLimiter(
"crawlStatus" as RateLimiterMode, // "crawlStatus" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"growth", // "growth",
); // );
expect(limiter4.points).toBe(250); // expect(limiter4.points).toBe(250);
}); // });
it("should return the default rate limiter if plan is not provided", () => { // it("should return the default rate limiter if plan is not provided", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"crawl" as RateLimiterMode, // "crawl" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
); // );
expect(limiter.points).toBe(3); // expect(limiter.points).toBe(3);
const limiter2 = getRateLimiter( // const limiter2 = getRateLimiter(
"scrape" as RateLimiterMode, // "scrape" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
); // );
expect(limiter2.points).toBe(20); // expect(limiter2.points).toBe(20);
}); // });
it("should create a new RateLimiterRedis instance with correct parameters", () => { // it("should create a new RateLimiterRedis instance with correct parameters", () => {
const keyPrefix = "test-prefix"; // const keyPrefix = "test-prefix";
const points = 10; // const points = 10;
const limiter = new RateLimiterRedis({ // const limiter = new RateLimiterRedis({
storeClient: redisRateLimitClient, // storeClient: redisRateLimitClient,
keyPrefix, // keyPrefix,
points, // points,
duration: 60, // duration: 60,
}); // });
expect(limiter.keyPrefix).toBe(keyPrefix); // expect(limiter.keyPrefix).toBe(keyPrefix);
expect(limiter.points).toBe(points); // expect(limiter.points).toBe(points);
expect(limiter.duration).toBe(60); // expect(limiter.duration).toBe(60);
}); // });
it("should return the correct rate limiter for 'preview' mode", () => { // it("should return the correct rate limiter for 'preview' mode", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"preview" as RateLimiterMode, // "preview" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"free", // "free",
); // );
expect(limiter.points).toBe(5); // expect(limiter.points).toBe(5);
const limiter2 = getRateLimiter( // const limiter2 = getRateLimiter(
"preview" as RateLimiterMode, // "preview" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
); // );
expect(limiter2.points).toBe(5); // expect(limiter2.points).toBe(5);
}); // });
it("should return the correct rate limiter for 'account' mode", () => { // it("should return the correct rate limiter for 'account' mode", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"account" as RateLimiterMode, // "account" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"free", // "free",
); // );
expect(limiter.points).toBe(100); // expect(limiter.points).toBe(100);
const limiter2 = getRateLimiter( // const limiter2 = getRateLimiter(
"account" as RateLimiterMode, // "account" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
); // );
expect(limiter2.points).toBe(100); // expect(limiter2.points).toBe(100);
}); // });
it("should return the correct rate limiter for 'crawlStatus' mode", () => { // it("should return the correct rate limiter for 'crawlStatus' mode", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"crawlStatus" as RateLimiterMode, // "crawlStatus" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"free", // "free",
); // );
expect(limiter.points).toBe(150); // expect(limiter.points).toBe(150);
const limiter2 = getRateLimiter( // const limiter2 = getRateLimiter(
"crawlStatus" as RateLimiterMode, // "crawlStatus" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
); // );
expect(limiter2.points).toBe(250); // expect(limiter2.points).toBe(250);
}); // });
it("should consume points correctly for 'crawl' mode", async () => { // it("should consume points correctly for 'crawl' mode", async () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"crawl" as RateLimiterMode, // "crawl" as RateLimiterMode,
"test-prefix:someTokenCRAWL", // "test-prefix:someTokenCRAWL",
"free", // "free",
); // );
const consumePoints = 1; // const consumePoints = 1;
const res = await limiter.consume( // const res = await limiter.consume(
"test-prefix:someTokenCRAWL", // "test-prefix:someTokenCRAWL",
consumePoints, // consumePoints,
); // );
expect(res.remainingPoints).toBe(1); // expect(res.remainingPoints).toBe(1);
}); // });
it("should consume points correctly for 'scrape' mode (DEFAULT)", async () => { // it("should consume points correctly for 'scrape' mode (DEFAULT)", async () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"scrape" as RateLimiterMode, // "scrape" as RateLimiterMode,
"test-prefix:someTokenX", // "test-prefix:someTokenX",
); // );
const consumePoints = 4; // const consumePoints = 4;
const res = await limiter.consume("test-prefix:someTokenX", consumePoints); // const res = await limiter.consume("test-prefix:someTokenX", consumePoints);
expect(res.remainingPoints).toBe(16); // expect(res.remainingPoints).toBe(16);
}); // });
it("should consume points correctly for 'scrape' mode (HOBBY)", async () => { // it("should consume points correctly for 'scrape' mode (HOBBY)", async () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"scrape" as RateLimiterMode, // "scrape" as RateLimiterMode,
"test-prefix:someTokenXY", // "test-prefix:someTokenXY",
"hobby", // "hobby",
); // );
expect(limiter.points).toBe(20); // expect(limiter.points).toBe(20);
const consumePoints = 5; // const consumePoints = 5;
const res = await limiter.consume("test-prefix:someTokenXY", consumePoints); // const res = await limiter.consume("test-prefix:someTokenXY", consumePoints);
expect(res.consumedPoints).toBe(5); // expect(res.consumedPoints).toBe(5);
expect(res.remainingPoints).toBe(15); // expect(res.remainingPoints).toBe(15);
}); // });
it("should return the correct rate limiter for 'crawl' mode", () => { // it("should return the correct rate limiter for 'crawl' mode", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"crawl" as RateLimiterMode, // "crawl" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"free", // "free",
); // );
expect(limiter.points).toBe(2); // expect(limiter.points).toBe(2);
const limiter2 = getRateLimiter( // const limiter2 = getRateLimiter(
"crawl" as RateLimiterMode, // "crawl" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"starter", // "starter",
); // );
expect(limiter2.points).toBe(10); // expect(limiter2.points).toBe(10);
const limiter3 = getRateLimiter( // const limiter3 = getRateLimiter(
"crawl" as RateLimiterMode, // "crawl" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"standard", // "standard",
); // );
expect(limiter3.points).toBe(5); // expect(limiter3.points).toBe(5);
}); // });
it("should return the correct rate limiter for 'scrape' mode", () => { // it("should return the correct rate limiter for 'scrape' mode", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"scrape" as RateLimiterMode, // "scrape" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"free", // "free",
); // );
expect(limiter.points).toBe(10); // expect(limiter.points).toBe(10);
const limiter2 = getRateLimiter( // const limiter2 = getRateLimiter(
"scrape" as RateLimiterMode, // "scrape" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"starter", // "starter",
); // );
expect(limiter2.points).toBe(100); // expect(limiter2.points).toBe(100);
const limiter3 = getRateLimiter( // const limiter3 = getRateLimiter(
"scrape" as RateLimiterMode, // "scrape" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"standard", // "standard",
); // );
expect(limiter3.points).toBe(100); // expect(limiter3.points).toBe(100);
const limiter4 = getRateLimiter( // const limiter4 = getRateLimiter(
"scrape" as RateLimiterMode, // "scrape" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"growth", // "growth",
); // );
expect(limiter4.points).toBe(1000); // expect(limiter4.points).toBe(1000);
}); // });
it("should return the correct rate limiter for 'search' mode", () => { // it("should return the correct rate limiter for 'search' mode", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"search" as RateLimiterMode, // "search" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"free", // "free",
); // );
expect(limiter.points).toBe(5); // expect(limiter.points).toBe(5);
const limiter2 = getRateLimiter( // const limiter2 = getRateLimiter(
"search" as RateLimiterMode, // "search" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"starter", // "starter",
); // );
expect(limiter2.points).toBe(50); // expect(limiter2.points).toBe(50);
const limiter3 = getRateLimiter( // const limiter3 = getRateLimiter(
"search" as RateLimiterMode, // "search" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"standard", // "standard",
); // );
expect(limiter3.points).toBe(50); // expect(limiter3.points).toBe(50);
}); // });
it("should return the correct rate limiter for 'preview' mode", () => { // it("should return the correct rate limiter for 'preview' mode", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"preview" as RateLimiterMode, // "preview" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"free", // "free",
); // );
expect(limiter.points).toBe(5); // expect(limiter.points).toBe(5);
const limiter2 = getRateLimiter( // const limiter2 = getRateLimiter(
"preview" as RateLimiterMode, // "preview" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
); // );
expect(limiter2.points).toBe(5); // expect(limiter2.points).toBe(5);
}); // });
it("should return the correct rate limiter for 'account' mode", () => { // it("should return the correct rate limiter for 'account' mode", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"account" as RateLimiterMode, // "account" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"free", // "free",
); // );
expect(limiter.points).toBe(100); // expect(limiter.points).toBe(100);
const limiter2 = getRateLimiter( // const limiter2 = getRateLimiter(
"account" as RateLimiterMode, // "account" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
); // );
expect(limiter2.points).toBe(100); // expect(limiter2.points).toBe(100);
}); // });
it("should return the correct rate limiter for 'crawlStatus' mode", () => { // it("should return the correct rate limiter for 'crawlStatus' mode", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"crawlStatus" as RateLimiterMode, // "crawlStatus" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"free", // "free",
); // );
expect(limiter.points).toBe(150); // expect(limiter.points).toBe(150);
const limiter2 = getRateLimiter( // const limiter2 = getRateLimiter(
"crawlStatus" as RateLimiterMode, // "crawlStatus" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
); // );
expect(limiter2.points).toBe(250); // expect(limiter2.points).toBe(250);
}); // });
it("should return the correct rate limiter for 'testSuite' mode", () => { // it("should return the correct rate limiter for 'testSuite' mode", () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"testSuite" as RateLimiterMode, // "testSuite" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
"free", // "free",
); // );
expect(limiter.points).toBe(10000); // expect(limiter.points).toBe(10000);
const limiter2 = getRateLimiter( // const limiter2 = getRateLimiter(
"testSuite" as RateLimiterMode, // "testSuite" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
); // );
expect(limiter2.points).toBe(10000); // expect(limiter2.points).toBe(10000);
}); // });
it("should throw an error when consuming more points than available", async () => { // it("should throw an error when consuming more points than available", async () => {
const limiter = getRateLimiter( // const limiter = getRateLimiter(
"crawl" as RateLimiterMode, // "crawl" as RateLimiterMode,
"test-prefix:someToken", // "test-prefix:someToken",
); // );
const consumePoints = limiter.points + 1; // const consumePoints = limiter.points + 1;
try { // try {
await limiter.consume("test-prefix:someToken", consumePoints); // await limiter.consume("test-prefix:someToken", consumePoints);
} catch (error) { // } catch (error) {
// expect remaining points to be 0 // // expect remaining points to be 0
const res = await limiter.get("test-prefix:someToken"); // const res = await limiter.get("test-prefix:someToken");
expect(res?.remainingPoints).toBe(0); // expect(res?.remainingPoints).toBe(0);
} // }
}); // });
it("should reset points after duration", async () => { // it("should reset points after duration", async () => {
const keyPrefix = "test-prefix"; // const keyPrefix = "test-prefix";
const points = 10; // const points = 10;
const duration = 1; // 1 second // const duration = 1; // 1 second
const limiter = new RateLimiterRedis({ // const limiter = new RateLimiterRedis({
storeClient: redisRateLimitClient, // storeClient: redisRateLimitClient,
keyPrefix, // keyPrefix,
points, // points,
duration, // duration,
}); // });
const consumePoints = 5; // const consumePoints = 5;
await limiter.consume("test-prefix:someToken", consumePoints); // await limiter.consume("test-prefix:someToken", consumePoints);
await new Promise((resolve) => setTimeout(resolve, duration * 1000 + 100)); // Wait for duration + 100ms // await new Promise((resolve) => setTimeout(resolve, duration * 1000 + 100)); // Wait for duration + 100ms
const res = await limiter.consume("test-prefix:someToken", consumePoints); // const res = await limiter.consume("test-prefix:someToken", consumePoints);
expect(res.remainingPoints).toBe(points - consumePoints); // expect(res.remainingPoints).toBe(points - consumePoints);
}); // });
}); // });
// TODO: FIX

307
apps/api/src/services/rate-limiter.ts
View File

@ -1,162 +1,7 @@
import { RateLimiterRedis } from "rate-limiter-flexible"; import { RateLimiterRedis } from "rate-limiter-flexible";
import { PlanType, RateLimiterMode } from "../../src/types"; import { RateLimiterMode } from "../types";
import Redis from "ioredis"; import Redis from "ioredis";
import type { AuthCreditUsageChunk } from "../controllers/v1/types";
export const CONCURRENCY_LIMIT: Omit<Record<PlanType, number>, ""> = {
free: 2,
hobby: 5,
starter: 50,
standard: 50,
standardNew: 50,
standardnew: 50,
scale: 100,
growth: 100,
growthdouble: 100,
etier2c: 300,
etier1a: 200,
etier2a: 300,
etierscale1: 150,
etierscale2: 200,
testSuite: 200,
devB: 120,
etier2d: 250,
manual: 200,
extract_starter: 20,
extract_explorer: 100,
extract_pro: 200
};
const RATE_LIMITS = {
crawl: {
default: 15,
free: 2,
starter: 50,
standard: 25,
standardOld: 200,
scale: 250,
hobby: 15,
standardNew: 50,
standardnew: 50,
growth: 250,
growthdouble: 250,
etier2c: 1500,
etier1a: 5000,
etier2a: 1500,
etierscale1: 750,
etierscale2: 1500,
// extract ops
extract_starter: 100,
extract_explorer: 500,
extract_pro: 1000,
},
scrape: {
default: 100,
free: 10,
starter: 500,
standard: 500,
standardOld: 500,
scale: 2500,
hobby: 100,
standardNew: 500,
standardnew: 500,
growth: 5000,
growthdouble: 5000,
etier2c: 12500,
etier1a: 5000,
etier2a: 12500,
etierscale1: 7500,
etierscale2: 12500,
// extract ops
extract_starter: 100,
extract_explorer: 500,
extract_pro: 1000,
},
search: {
default: 100,
free: 5,
starter: 250,
standard: 250,
standardOld: 200,
scale: 2500,
hobby: 50,
standardNew: 250,
standardnew: 250,
growth: 2500,
growthdouble: 2500,
etier2c: 12500,
etier1a: 5000,
etier2a: 12500,
etierscale1: 7500,
etierscale2: 12500,
// extract ops
extract_starter: 100,
extract_explorer: 500,
extract_pro: 1000,
},
map: {
default: 100,
free: 5,
starter: 250,
standard: 250,
standardOld: 250,
scale: 2500,
hobby: 50,
standardNew: 250,
standardnew: 250,
growth: 5000,
growthdouble: 5000,
etier2c: 12500,
etier1a: 5000,
etier2a: 12500,
etierscale1: 7500,
etierscale2: 12500,
// extract ops
extract_starter: 100,
extract_explorer: 500,
extract_pro: 1000,
},
extract: {
default: 100,
free: 10,
starter: 500,
standard: 500,
standardOld: 500,
scale: 1000,
hobby: 100,
standardNew: 500,
standardnew: 500,
growth: 1000,
growthdouble: 1000,
etier2c: 1000,
etier1a: 1000,
etier2a: 1000,
etierscale1: 1000,
etierscale2: 1000,
extract_starter: 100,
extract_explorer: 500,
extract_pro: 1000,
},
preview: {
free: 5,
default: 25,
},
account: {
free: 100,
default: 500,
},
crawlStatus: {
free: 500,
default: 25000,
},
extractStatus: {
free: 500,
default: 25000,
},
testSuite: {
free: 10000,
default: 50000,
},
};
export const redisRateLimitClient = new Redis( export const redisRateLimitClient = new Redis(
process.env.REDIS_RATE_LIMIT_URL!, process.env.REDIS_RATE_LIMIT_URL!,
@ -170,11 +15,6 @@ const createRateLimiter = (keyPrefix, points) =>
duration: 60, // Duration in seconds duration: 60, // Duration in seconds
}); });
export const serverRateLimiter = createRateLimiter(
"server",
RATE_LIMITS.account.default,
);
export const testSuiteRateLimiter = new RateLimiterRedis({ export const testSuiteRateLimiter = new RateLimiterRedis({
storeClient: redisRateLimitClient, storeClient: redisRateLimitClient,
keyPrefix: "test-suite", keyPrefix: "test-suite",
@ -182,41 +22,7 @@ export const testSuiteRateLimiter = new RateLimiterRedis({
duration: 60, // Duration in seconds duration: 60, // Duration in seconds
}); });
export const devBRateLimiter = new RateLimiterRedis({ // TODO: PUT OVERRIDES FOR THESE INTO THE DB - mogery
storeClient: redisRateLimitClient,
keyPrefix: "dev-b",
points: 1200,
duration: 60, // Duration in seconds
});
export const manualRateLimiter = new RateLimiterRedis({
storeClient: redisRateLimitClient,
keyPrefix: "manual",
points: 10000,
duration: 60, // Duration in seconds
});
export const scrapeStatusRateLimiter = new RateLimiterRedis({
storeClient: redisRateLimitClient,
keyPrefix: "scrape-status",
points: 400,
duration: 60, // Duration in seconds
});
export const etier1aRateLimiter = new RateLimiterRedis({
storeClient: redisRateLimitClient,
keyPrefix: "etier1a",
points: 10000,
duration: 60, // Duration in seconds
});
export const etier2aRateLimiter = new RateLimiterRedis({
storeClient: redisRateLimitClient,
keyPrefix: "etier2a",
points: 2500,
duration: 60, // Duration in seconds
});
const testSuiteTokens = [ const testSuiteTokens = [
"a01ccae", "a01ccae",
"6254cf9", "6254cf9",
@ -240,105 +46,32 @@ const testSuiteTokens = [
"0a18c9e", // gh "0a18c9e", // gh
]; ];
const manual_growth = ["22a07b64-cbfe-4924-9273-e3f01709cdf2"]; // TODO: PUT OVERRIDES FOR THESE INTO THE DB - mogery
const manual = ["69be9e74-7624-4990-b20d-08e0acc70cf6", "9661a311-3d75-45d2-bb70-71004d995873"]; // const manual_growth = ["22a07b64-cbfe-4924-9273-e3f01709cdf2"];
const manual_etier2c = ["77545e01-9cec-4fa9-8356-883fc66ac13e", "778c62c4-306f-4039-b372-eb20174760c0"]; // const manual = ["69be9e74-7624-4990-b20d-08e0acc70cf6", "9661a311-3d75-45d2-bb70-71004d995873"];
// const manual_etier2c = ["77545e01-9cec-4fa9-8356-883fc66ac13e", "778c62c4-306f-4039-b372-eb20174760c0"];
function makePlanKey(plan?: string) { const fallbackRateLimits: AuthCreditUsageChunk["rate_limits"] = {
return plan ? plan.replace("-", "") : "default"; // "default" crawl: 15,
} scrape: 100,
search: 100,
export function getRateLimiterPoints( map: 100,
mode: RateLimiterMode, extract: 100,
token?: string, preview: 25,
plan?: string, extractStatus: 25000,
teamId?: string, crawlStatus: 25000,
): number { };
const rateLimitConfig = RATE_LIMITS[mode]; // {default : 5}
if (!rateLimitConfig) return RATE_LIMITS.account.default;
const points: number =
rateLimitConfig[makePlanKey(plan)] || rateLimitConfig.default; // 5
return points;
}
export function getRateLimiter( export function getRateLimiter(
mode: RateLimiterMode, mode: RateLimiterMode,
token?: string, rate_limits: AuthCreditUsageChunk["rate_limits"] | null,
plan?: string,
teamId?: string,
): RateLimiterRedis { ): RateLimiterRedis {
if (token && testSuiteTokens.some((testToken) => token.includes(testToken))) {
return testSuiteRateLimiter;
}
if (teamId && teamId === process.env.DEV_B_TEAM_ID) {
return devBRateLimiter;
}
if (teamId && (teamId === process.env.ETIER1A_TEAM_ID || teamId === process.env.ETIER1A_TEAM_ID_O)) {
return etier1aRateLimiter;
}
if (teamId && (teamId === process.env.ETIER2A_TEAM_ID || teamId === process.env.ETIER2A_TEAM_ID_B)) {
return etier2aRateLimiter;
}
if (teamId && teamId === process.env.ETIER2D_TEAM_ID) {
return etier2aRateLimiter;
}
if (teamId && (manual.includes(teamId) || manual_etier2c.includes(teamId))) {
return manualRateLimiter;
}
return createRateLimiter( return createRateLimiter(
`${mode}-${makePlanKey(plan)}`, `${mode}`,
getRateLimiterPoints(mode, token, plan, teamId), (rate_limits ?? fallbackRateLimits)[mode] ?? 500,
); );
} }
export function getConcurrencyLimitMax(
plan: PlanType,
teamId?: string,
): number {
// Moved this to auth check, plan will come as testSuite if token is present
// if (token && testSuiteTokens.some((testToken) => token.includes(testToken))) {
// return CONCURRENCY_LIMIT.testSuite;
// }
if (teamId && teamId === process.env.DEV_B_TEAM_ID) {
return CONCURRENCY_LIMIT.devB;
}
if (teamId && (teamId === process.env.ETIER1A_TEAM_ID || teamId === process.env.ETIER1A_TEAM_ID_O)) {
return CONCURRENCY_LIMIT.etier1a;
}
if (teamId && (teamId === process.env.ETIER2A_TEAM_ID || teamId === process.env.ETIER2A_TEAM_ID_B)) {
return CONCURRENCY_LIMIT.etier2a;
}
if (teamId && teamId === process.env.ETIER2D_TEAM_ID) {
return CONCURRENCY_LIMIT.etier2a;
}
if (teamId && manual.includes(teamId)) {
return CONCURRENCY_LIMIT.manual;
}
if (teamId && manual_etier2c.includes(teamId)) {
return CONCURRENCY_LIMIT.etier2c;
}
if (teamId && manual_growth.includes(teamId)) {
return CONCURRENCY_LIMIT.growth;
}
return CONCURRENCY_LIMIT[plan] ?? 10;
}
export function isTestSuiteToken(token: string): boolean { export function isTestSuiteToken(token: string): boolean {
return testSuiteTokens.some((testToken) => token.includes(testToken)); return testSuiteTokens.some((testToken) => token.includes(testToken));
} }

26
apps/api/src/types.ts
View File

@ -37,7 +37,6 @@ export interface WebScraperOptions {
scrapeOptions: ScrapeOptions; scrapeOptions: ScrapeOptions;
internalOptions?: InternalOptions; internalOptions?: InternalOptions;
team_id: string; team_id: string;
plan: string;
origin?: string; origin?: string;
crawl_id?: string; crawl_id?: string;
sitemapped?: boolean; sitemapped?: boolean;
@ -144,7 +143,6 @@ export type AuthResponse =
success: true; success: true;
team_id: string; team_id: string;
api_key?: string; api_key?: string;
plan?: PlanType;
chunk: AuthCreditUsageChunk | null; chunk: AuthCreditUsageChunk | null;
} }
| { | {
@ -178,30 +176,6 @@ export type ScrapeLog = {
ipv6_support?: boolean | null; ipv6_support?: boolean | null;
}; };
export type PlanType =
| "starter"
| "standard"
| "scale"
| "hobby"
| "standardnew"
| "standardNew"
| "growth"
| "growthdouble"
| "etier2c"
| "etier1a"
| "etierscale1"
| "etierscale2"
| "etier2a"
| "free"
| "testSuite"
| "devB"
| "etier2d"
| "manual"
| "extract_starter"
| "extract_explorer"
| "extract_pro"
| "";
export type WebhookEventType = export type WebhookEventType =
| "crawl.page" | "crawl.page"
| "batch_scrape.page" | "batch_scrape.page"