diff --git a/apps/api/src/controllers/v0/crawl-cancel.ts b/apps/api/src/controllers/v0/crawl-cancel.ts index db834230..16f6cc87 100644 --- a/apps/api/src/controllers/v0/crawl-cancel.ts +++ b/apps/api/src/controllers/v0/crawl-cancel.ts @@ -6,6 +6,7 @@ import { logger } from "../../../src/lib/logger"; import { getCrawl, saveCrawl } from "../../../src/lib/crawl-redis"; import * as Sentry from "@sentry/node"; import { configDotenv } from "dotenv"; +import { redisConnection } from "../../services/queue-service"; configDotenv(); export async function crawlCancelController(req: Request, res: Response) { @@ -19,6 +20,9 @@ export async function crawlCancelController(req: Request, res: Response) { const { team_id } = auth; + redisConnection.sadd("teams_using_v0", team_id) + .catch(error => logger.error("Failed to add team to teams_using_v0", { error, team_id })); + const sc = await getCrawl(req.params.jobId); if (!sc) { return res.status(404).json({ error: "Job not found" }); diff --git a/apps/api/src/controllers/v0/crawl-status.ts b/apps/api/src/controllers/v0/crawl-status.ts index 7040a857..26ee0ee9 100644 --- a/apps/api/src/controllers/v0/crawl-status.ts +++ b/apps/api/src/controllers/v0/crawl-status.ts @@ -1,7 +1,7 @@ import { Request, Response } from "express"; import { authenticateUser } from "../auth"; import { RateLimiterMode } from "../../../src/types"; -import { getScrapeQueue } from "../../../src/services/queue-service"; +import { getScrapeQueue, redisConnection } from "../../../src/services/queue-service"; import { logger } from "../../../src/lib/logger"; import { getCrawl, getCrawlJobs } from "../../../src/lib/crawl-redis"; import { supabaseGetJobsByCrawlId } from "../../../src/lib/supabase-jobs"; @@ -67,6 +67,9 @@ export async function crawlStatusController(req: Request, res: Response) { const { team_id } = auth; + redisConnection.sadd("teams_using_v0", team_id) + .catch(error => logger.error("Failed to add team to teams_using_v0", { error, team_id })); + const sc = await getCrawl(req.params.jobId); if (!sc) { return res.status(404).json({ error: "Job not found" }); diff --git a/apps/api/src/controllers/v0/crawl.ts b/apps/api/src/controllers/v0/crawl.ts index 6b15a219..2eba651d 100644 --- a/apps/api/src/controllers/v0/crawl.ts +++ b/apps/api/src/controllers/v0/crawl.ts @@ -24,7 +24,7 @@ import { saveCrawl, StoredCrawl, } from "../../../src/lib/crawl-redis"; -import { getScrapeQueue } from "../../../src/services/queue-service"; +import { getScrapeQueue, redisConnection } from "../../../src/services/queue-service"; import { checkAndUpdateURL } from "../../../src/lib/validateUrl"; import * as Sentry from "@sentry/node"; import { getJobPriority } from "../../lib/job-priority"; @@ -41,6 +41,9 @@ export async function crawlController(req: Request, res: Response) { const { team_id, plan, chunk } = auth; + redisConnection.sadd("teams_using_v0", team_id) + .catch(error => logger.error("Failed to add team to teams_using_v0", { error, team_id })); + if (req.headers["x-idempotency-key"]) { const isIdempotencyValid = await validateIdempotencyKey(req); if (!isIdempotencyValid) { diff --git a/apps/api/src/controllers/v0/keyAuth.ts b/apps/api/src/controllers/v0/keyAuth.ts index 2495705c..ce8dc7ce 100644 --- a/apps/api/src/controllers/v0/keyAuth.ts +++ b/apps/api/src/controllers/v0/keyAuth.ts @@ -2,6 +2,8 @@ import { AuthResponse, RateLimiterMode } from "../../types"; import { Request, Response } from "express"; import { authenticateUser } from "../auth"; +import { redisConnection } from "../../services/queue-service"; +import { logger } from "../../lib/logger"; export const keyAuthController = async (req: Request, res: Response) => { try { @@ -11,6 +13,9 @@ export const keyAuthController = async (req: Request, res: Response) => { return res.status(auth.status).json({ error: auth.error }); } + redisConnection.sadd("teams_using_v0", auth.team_id) + .catch(error => logger.error("Failed to add team to teams_using_v0", { error, team_id: auth.team_id })); + // if success, return success: true return res.status(200).json({ success: true }); } catch (error) { diff --git a/apps/api/src/controllers/v0/scrape.ts b/apps/api/src/controllers/v0/scrape.ts index 05bf364b..62d62b09 100644 --- a/apps/api/src/controllers/v0/scrape.ts +++ b/apps/api/src/controllers/v0/scrape.ts @@ -21,7 +21,7 @@ import { defaultOrigin, } from "../../lib/default-values"; import { addScrapeJob, waitForJob } from "../../services/queue-jobs"; -import { getScrapeQueue } from "../../services/queue-service"; +import { getScrapeQueue, redisConnection } from "../../services/queue-service"; import { v4 as uuidv4 } from "uuid"; import { logger } from "../../lib/logger"; import * as Sentry from "@sentry/node"; @@ -181,6 +181,9 @@ export async function scrapeController(req: Request, res: Response) { const { team_id, plan, chunk } = auth; + redisConnection.sadd("teams_using_v0", team_id) + .catch(error => logger.error("Failed to add team to teams_using_v0", { error, team_id })); + const crawlerOptions = req.body.crawlerOptions ?? {}; const pageOptions = { ...defaultPageOptions, ...req.body.pageOptions }; const extractorOptions = { diff --git a/apps/api/src/controllers/v0/search.ts b/apps/api/src/controllers/v0/search.ts index 6a3513df..ac7d7f62 100644 --- a/apps/api/src/controllers/v0/search.ts +++ b/apps/api/src/controllers/v0/search.ts @@ -11,7 +11,7 @@ import { search } from "../../search"; import { isUrlBlocked } from "../../scraper/WebScraper/utils/blocklist"; import { v4 as uuidv4 } from "uuid"; import { logger } from "../../lib/logger"; -import { getScrapeQueue } from "../../services/queue-service"; +import { getScrapeQueue, redisConnection } from "../../services/queue-service"; import { addScrapeJob, waitForJob } from "../../services/queue-jobs"; import * as Sentry from "@sentry/node"; import { getJobPriority } from "../../lib/job-priority"; @@ -163,6 +163,10 @@ export async function searchController(req: Request, res: Response) { return res.status(auth.status).json({ error: auth.error }); } const { team_id, plan, chunk } = auth; + + redisConnection.sadd("teams_using_v0", team_id) + .catch(error => logger.error("Failed to add team to teams_using_v0", { error, team_id })); + const crawlerOptions = req.body.crawlerOptions ?? {}; const pageOptions = req.body.pageOptions ?? { includeHtml: req.body.pageOptions?.includeHtml ?? false,