mirror of
https://git.mirrors.martin98.com/https://github.com/mendableai/firecrawl
synced 2025-08-06 04:47:29 +08:00
Nick: fixed more stuff
This commit is contained in:
parent
c0541cc990
commit
18f9cd09e1
25
apps/api/src/controllers/v1/concurrency-check.ts
Normal file
25
apps/api/src/controllers/v1/concurrency-check.ts
Normal file
@ -0,0 +1,25 @@
|
||||
import { authenticateUser } from "../auth";
|
||||
import {
|
||||
ConcurrencyCheckParams,
|
||||
ConcurrencyCheckResponse,
|
||||
RequestWithAuth,
|
||||
} from "./types";
|
||||
import { RateLimiterMode } from "../../types";
|
||||
import { Response } from "express";
|
||||
import { redisConnection } from "../../services/queue-service";
|
||||
// Basically just middleware and error wrapping
|
||||
export async function concurrencyCheckController(
|
||||
req: RequestWithAuth<ConcurrencyCheckParams, undefined, undefined>,
|
||||
res: Response<ConcurrencyCheckResponse>
|
||||
) {
|
||||
const concurrencyLimiterKey = "concurrency-limiter:" + req.params.teamId;
|
||||
const now = Date.now();
|
||||
const activeJobsOfTeam = await redisConnection.zrangebyscore(
|
||||
concurrencyLimiterKey,
|
||||
now,
|
||||
Infinity
|
||||
);
|
||||
return res
|
||||
.status(200)
|
||||
.json({ success: true, concurrency: activeJobsOfTeam.length });
|
||||
}
|
@ -294,6 +294,17 @@ export type CrawlStatusParams = {
|
||||
jobId: string;
|
||||
};
|
||||
|
||||
export type ConcurrencyCheckParams = {
|
||||
teamId: string;
|
||||
};
|
||||
|
||||
export type ConcurrencyCheckResponse =
|
||||
| ErrorResponse
|
||||
| {
|
||||
success: true;
|
||||
concurrency: number;
|
||||
};
|
||||
|
||||
export type CrawlStatusResponse =
|
||||
| ErrorResponse
|
||||
| {
|
||||
|
@ -16,6 +16,7 @@ import { isUrlBlocked } from "../scraper/WebScraper/utils/blocklist";
|
||||
import { crawlCancelController } from "../controllers/v1/crawl-cancel";
|
||||
import { Logger } from "../lib/logger";
|
||||
import { scrapeStatusController } from "../controllers/v1/scrape-status";
|
||||
import { concurrencyCheckController } from "../controllers/v1/concurrency-check";
|
||||
// import { crawlPreviewController } from "../../src/controllers/v1/crawlPreview";
|
||||
// import { crawlJobStatusPreviewController } from "../../src/controllers/v1/status";
|
||||
// import { searchController } from "../../src/controllers/v1/search";
|
||||
@ -140,11 +141,19 @@ v1Router.get(
|
||||
wrap(scrapeStatusController)
|
||||
);
|
||||
|
||||
v1Router.get(
|
||||
"/concurrency-check",
|
||||
authMiddleware(RateLimiterMode.CrawlStatus),
|
||||
wrap(concurrencyCheckController)
|
||||
);
|
||||
|
||||
v1Router.ws(
|
||||
"/crawl/:jobId",
|
||||
crawlStatusWSController
|
||||
);
|
||||
|
||||
|
||||
|
||||
// v1Router.post("/crawlWebsitePreview", crawlPreviewController);
|
||||
|
||||
|
||||
|
@ -49,7 +49,7 @@ describe("Rate Limiter Service", () => {
|
||||
"nonexistent" as RateLimiterMode,
|
||||
"test-prefix:someToken"
|
||||
);
|
||||
expect(limiter).toBe(serverRateLimiter);
|
||||
expect(limiter.points).toBe(serverRateLimiter.points);
|
||||
});
|
||||
|
||||
it("should return the correct rate limiter based on mode and plan", () => {
|
||||
@ -210,7 +210,7 @@ describe("Rate Limiter Service", () => {
|
||||
"test-prefix:someToken",
|
||||
"starter"
|
||||
);
|
||||
expect(limiter2.points).toBe(3);
|
||||
expect(limiter2.points).toBe(10);
|
||||
|
||||
const limiter3 = getRateLimiter(
|
||||
"crawl" as RateLimiterMode,
|
||||
@ -233,7 +233,7 @@ describe("Rate Limiter Service", () => {
|
||||
"test-prefix:someToken",
|
||||
"starter"
|
||||
);
|
||||
expect(limiter2.points).toBe(20);
|
||||
expect(limiter2.points).toBe(100);
|
||||
|
||||
const limiter3 = getRateLimiter(
|
||||
"scrape" as RateLimiterMode,
|
||||
@ -263,14 +263,14 @@ describe("Rate Limiter Service", () => {
|
||||
"test-prefix:someToken",
|
||||
"starter"
|
||||
);
|
||||
expect(limiter2.points).toBe(20);
|
||||
expect(limiter2.points).toBe(50);
|
||||
|
||||
const limiter3 = getRateLimiter(
|
||||
"search" as RateLimiterMode,
|
||||
"test-prefix:someToken",
|
||||
"standard"
|
||||
);
|
||||
expect(limiter3.points).toBe(40);
|
||||
expect(limiter3.points).toBe(50);
|
||||
});
|
||||
|
||||
it("should return the correct rate limiter for 'preview' mode", () => {
|
||||
|
@ -132,8 +132,22 @@ export function getRateLimiterPoints(
|
||||
token?: string,
|
||||
plan?: string,
|
||||
teamId?: string
|
||||
) {
|
||||
) : 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(
|
||||
mode: RateLimiterMode,
|
||||
token?: string,
|
||||
plan?: string,
|
||||
teamId?: string
|
||||
) : RateLimiterRedis {
|
||||
if (token && testSuiteTokens.some(testToken => token.includes(testToken))) {
|
||||
return testSuiteRateLimiter;
|
||||
}
|
||||
@ -145,22 +159,6 @@ export function getRateLimiterPoints(
|
||||
if(teamId && manual.includes(teamId)) {
|
||||
return manualRateLimiter;
|
||||
}
|
||||
|
||||
const rateLimitConfig = RATE_LIMITS[mode]; // {default : 5}
|
||||
|
||||
if (!rateLimitConfig) return serverRateLimiter;
|
||||
|
||||
const points =
|
||||
rateLimitConfig[makePlanKey(plan)] || rateLimitConfig.default || rateLimitConfig; // 5
|
||||
|
||||
return points;
|
||||
}
|
||||
|
||||
export function getRateLimiter(
|
||||
mode: RateLimiterMode,
|
||||
token?: string,
|
||||
plan?: string,
|
||||
teamId?: string
|
||||
) {
|
||||
|
||||
return createRateLimiter(`${mode}-${makePlanKey(plan)}`, getRateLimiterPoints(mode, token, plan, teamId));
|
||||
}
|
||||
|
Loading…
x
Reference in New Issue
Block a user