diff --git a/apps/api/openapi.json b/apps/api/openapi.json index 98cacd8175..76809d30bd 100644 --- a/apps/api/openapi.json +++ b/apps/api/openapi.json @@ -2277,6 +2277,7 @@ "markdown", "html", "rawHtml", + "rawBase64", "links", "screenshot", "screenshot@fullPage", @@ -2285,7 +2286,7 @@ "branding" ] }, - "description": "Formats to include in the output.", + "description": "Formats to include in the output. `rawBase64` must be requested by itself.", "default": ["markdown"] }, "onlyMainContent": { @@ -2606,6 +2607,11 @@ "nullable": true, "description": "Raw HTML content of the page if `rawHtml` is in `formats`" }, + "rawBase64": { + "type": "string", + "nullable": true, + "description": "Base64-encoded original response body if `rawBase64` is in `formats`" + }, "screenshot": { "type": "string", "nullable": true, @@ -2866,6 +2872,11 @@ "nullable": true, "description": "Raw HTML content of the page if `includeRawHtml` is true" }, + "rawBase64": { + "type": "string", + "nullable": true, + "description": "Base64-encoded original response body if `rawBase64` is in `formats`" + }, "links": { "type": "array", "items": { @@ -3003,6 +3014,11 @@ "nullable": true, "description": "Raw HTML content of the page if `includeRawHtml` is true" }, + "rawBase64": { + "type": "string", + "nullable": true, + "description": "Base64-encoded original response body if `rawBase64` is in `formats`" + }, "links": { "type": "array", "items": { diff --git a/apps/api/src/__tests__/snips/mocks/raw-base64.json b/apps/api/src/__tests__/snips/mocks/raw-base64.json new file mode 100644 index 0000000000..c80ae3b13b --- /dev/null +++ b/apps/api/src/__tests__/snips/mocks/raw-base64.json @@ -0,0 +1,70 @@ +[ + { + "time": 1, + "options": { + "url": "/scrape", + "method": "POST", + "body": { + "url": "https://example.com/raw", + "engine": "chrome-cdp", + "format": "rawBase64" + } + }, + "result": { + "status": 200, + "headers": {}, + "body": "{\"timeTaken\":0.1,\"content\":\"\",\"url\":\"https://example.com/raw\",\"pageStatusCode\":200,\"responseHeaders\":{\"content-type\":\"text/html; charset=utf-8\"},\"file\":{\"name\":\"raw.html\",\"content\":\"PGh0bWw+cmF3PC9odG1sPg==\"}}" + } + }, + { + "time": 2, + "options": { + "url": "/scrape", + "method": "POST", + "body": { + "url": "https://example.com/raw.pdf", + "engine": "chrome-cdp", + "format": "rawBase64" + } + }, + "result": { + "status": 200, + "headers": {}, + "body": "{\"timeTaken\":0.1,\"content\":\"\",\"url\":\"https://example.com/raw.pdf\",\"pageStatusCode\":200,\"responseHeaders\":{\"content-type\":\"application/pdf\"},\"file\":{\"name\":\"raw.pdf\",\"content\":\"PGh0bWw+cmF3PC9odG1sPg==\"}}" + } + }, + { + "time": 3, + "options": { + "url": "/scrape", + "method": "POST", + "body": { + "url": "https://example.com/raw.docx", + "engine": "chrome-cdp", + "format": "rawBase64" + } + }, + "result": { + "status": 200, + "headers": {}, + "body": "{\"timeTaken\":0.1,\"content\":\"\",\"url\":\"https://example.com/raw.docx\",\"pageStatusCode\":200,\"responseHeaders\":{\"content-type\":\"application/vnd.openxmlformats-officedocument.wordprocessingml.document\"},\"file\":{\"name\":\"raw.docx\",\"content\":\"PGh0bWw+cmF3PC9odG1sPg==\"}}" + } + }, + { + "time": 4, + "options": { + "url": "/scrape", + "method": "POST", + "body": { + "url": "https://example.com/raw-error", + "engine": "chrome-cdp", + "format": "rawBase64" + } + }, + "result": { + "status": 200, + "headers": {}, + "body": "{\"timeTaken\":0.1,\"content\":\"Not found\",\"url\":\"https://example.com/raw-error\",\"pageStatusCode\":404,\"responseHeaders\":{\"content-type\":\"text/html; charset=utf-8\"}}" + } + } +] diff --git a/apps/api/src/__tests__/snips/v1/types-validation.test.ts b/apps/api/src/__tests__/snips/v1/types-validation.test.ts index 79d7639313..8831f958b2 100644 --- a/apps/api/src/__tests__/snips/v1/types-validation.test.ts +++ b/apps/api/src/__tests__/snips/v1/types-validation.test.ts @@ -59,6 +59,22 @@ describe("V1 Types Validation", () => { expect(result.timeout).toBe(60000); }); + it("should only allow rawBase64 as the sole format", () => { + expect( + scrapeRequestSchema.parse({ + url: "https://example.com/file", + formats: ["rawBase64"], + }).formats, + ).toEqual(["rawBase64"]); + + expect(() => + scrapeRequestSchema.parse({ + url: "https://example.com/file", + formats: ["markdown", "rawBase64"], + }), + ).toThrow("The rawBase64 format cannot be combined with other formats"); + }); + it("should reject invalid URL", () => { const input = { url: "not-a-url", diff --git a/apps/api/src/__tests__/snips/v2/types-validation.test.ts b/apps/api/src/__tests__/snips/v2/types-validation.test.ts index da202c891e..113747c949 100644 --- a/apps/api/src/__tests__/snips/v2/types-validation.test.ts +++ b/apps/api/src/__tests__/snips/v2/types-validation.test.ts @@ -1,6 +1,7 @@ import { z } from "zod"; import { scrapeRequestSchema, + parseRequestSchema, scrapeOptions, extractRequestSchema, crawlRequestSchema, @@ -83,6 +84,22 @@ describe("V2 Types Validation", () => { expect(result.formats).toEqual([{ type: "markdown" }, { type: "html" }]); }); + it("should only allow rawBase64 as the sole format", () => { + expect( + scrapeRequestSchema.parse({ + url: "https://example.com/file", + formats: ["rawBase64"], + }).formats, + ).toEqual([{ type: "rawBase64" }]); + + expect(() => + scrapeRequestSchema.parse({ + url: "https://example.com/file", + formats: ["markdown", "rawBase64"], + }), + ).toThrow("The rawBase64 format cannot be combined with other formats"); + }); + it("should accept video format as string and object", () => { const stringInput: ScrapeRequestInput = { url: "https://example.com", @@ -685,6 +702,22 @@ describe("V2 Types Validation", () => { }); }); + describe("parseRequestSchema", () => { + it("should reject rawBase64 for file uploads", () => { + expect(() => + parseRequestSchema.parse({ + formats: ["rawBase64"], + file: { + buffer: Buffer.from("raw upload"), + filename: "upload.html", + contentType: "text/html", + kind: "html", + }, + }), + ).toThrow("The rawBase64 format is not supported for parse uploads"); + }); + }); + describe("extractRequestSchema", () => { it("should accept valid extract request with urls", () => { const input: ExtractRequestInput = { diff --git a/apps/api/src/controllers/__tests__/auth.test.ts b/apps/api/src/controllers/__tests__/auth.test.ts index 1f9b80d855..10bcb22efa 100644 --- a/apps/api/src/controllers/__tests__/auth.test.ts +++ b/apps/api/src/controllers/__tests__/auth.test.ts @@ -51,11 +51,24 @@ vi.mock("../../db/rpc", () => ({ authCreditUsageChunkFromTeam: vi.fn(), })); -vi.mock("../../services/rate-limiter", () => ({ - getRateLimiter: vi.fn(), - getAutumnRateLimiter: vi.fn(), +// The limiter builders are mocked, but getRateLimitOverride is kept real: it is +// the single source of truth for override resolution, and auth.ts calls it to +// decide whether the Autumn multiplier is needed at all. Stub ioredis so +// importing the real module doesn't open a connection. +vi.mock("ioredis", () => ({ + default: class {}, })); +vi.mock("../../services/rate-limiter", async importOriginal => { + const actual = + await importOriginal(); + return { + ...actual, + getRateLimiter: vi.fn(), + getAutumnRateLimiter: vi.fn(), + }; +}); + vi.mock("../../lib/keyless", async importOriginal => { const actual = await importOriginal(); return { @@ -88,6 +101,7 @@ describe("authenticateUser", () => { config.MCP_DELEGATED_CREDENTIAL_SECRET; const originalIntrospectUrl = config.OAUTH_INTROSPECT_URL; const originalIntrospectSecret = config.OAUTH_INTROSPECT_SECRET; + const originalPreviewToken = config.PREVIEW_TOKEN; beforeEach(() => { vi.mocked(isKeylessConfigured).mockReturnValue(false); @@ -105,6 +119,7 @@ describe("authenticateUser", () => { originalMcpDelegatedCredentialSecret; config.OAUTH_INTROSPECT_URL = originalIntrospectUrl; config.OAUTH_INTROSPECT_SECRET = originalIntrospectSecret; + config.PREVIEW_TOKEN = originalPreviewToken; vi.unstubAllGlobals(); vi.clearAllMocks(); }); @@ -497,6 +512,105 @@ describe("authenticateUser", () => { }); }); + it("passes the org rate-limit overrides to the API-key rate limiter", async () => { + config.USE_DB_AUTHENTICATION = true; + vi.mocked(getValue).mockResolvedValue(null); + const flags = { rateLimitOverrides: { scrape: 42 } }; + vi.mocked(authCreditUsageChunk).mockResolvedValue([ + { + api_key: "00000000-0000-4000-8000-000000000000", + api_key_id: 1, + team_id: "team-1", + org_id: "org-1", + flags, + }, + ]); + vi.mocked(redlock.using).mockImplementation( + async (_keys, _ttl, _options, fn) => fn({ aborted: false } as never), + ); + vi.mocked(autumnService.getRateLimitMultiplier).mockResolvedValue(50); + + const auth = await authenticateUser( + { + headers: { + authorization: "Bearer 00000000-0000-4000-8000-000000000000", + }, + socket: { remoteAddress: "127.0.0.1" }, + }, + {}, + RateLimiterMode.Scrape, + ); + + expect(auth.success).toBe(true); + // The override replaces the whole base × multiplier computation, so the + // Autumn multiplier is never fetched and a neutral 1 is passed instead. + expect(autumnService.getRateLimitMultiplier).not.toHaveBeenCalled(); + expect(getAutumnRateLimiter).toHaveBeenCalledWith( + RateLimiterMode.Scrape, + 1, + flags, + ); + }); + + it("still fetches the Autumn multiplier when no override covers the mode", async () => { + config.USE_DB_AUTHENTICATION = true; + vi.mocked(getValue).mockResolvedValue(null); + const flags = { rateLimitOverrides: { crawl: 42 } }; + vi.mocked(authCreditUsageChunk).mockResolvedValue([ + { + api_key: "00000000-0000-4000-8000-000000000000", + api_key_id: 1, + team_id: "team-1", + org_id: "org-1", + flags, + }, + ]); + vi.mocked(redlock.using).mockImplementation( + async (_keys, _ttl, _options, fn) => fn({ aborted: false } as never), + ); + vi.mocked(autumnService.getRateLimitMultiplier).mockResolvedValue(50); + + const auth = await authenticateUser( + { + headers: { + authorization: "Bearer 00000000-0000-4000-8000-000000000000", + }, + socket: { remoteAddress: "127.0.0.1" }, + }, + {}, + RateLimiterMode.Scrape, + ); + + expect(auth.success).toBe(true); + expect(autumnService.getRateLimitMultiplier).toHaveBeenCalledTimes(1); + expect(getAutumnRateLimiter).toHaveBeenCalledWith( + RateLimiterMode.Scrape, + 50, + flags, + ); + }); + + it("leaves the preview token on the static rate limiter", async () => { + config.USE_DB_AUTHENTICATION = true; + config.PREVIEW_TOKEN = "preview-token"; + vi.mocked(getRateLimiter).mockReturnValue({ + consume: vi.fn().mockResolvedValue(undefined), + } as never); + + const auth = await authenticateUser( + { + headers: { authorization: "Bearer preview-token" }, + socket: { remoteAddress: "127.0.0.1" }, + }, + {}, + RateLimiterMode.Scrape, + ); + + expect(auth.success).toBe(true); + expect(getRateLimiter).toHaveBeenCalledWith(RateLimiterMode.Preview); + expect(getAutumnRateLimiter).not.toHaveBeenCalled(); + }); + it("clears purpose-qualified and legacy ACUC cache entries", async () => { await clearACUC("api-key"); diff --git a/apps/api/src/controllers/auth.ts b/apps/api/src/controllers/auth.ts index 0c1bb4aad5..780a426afb 100644 --- a/apps/api/src/controllers/auth.ts +++ b/apps/api/src/controllers/auth.ts @@ -5,7 +5,11 @@ import { logger } from "../lib/logger"; import { parseApi } from "../lib/parseApi"; import { withAuth } from "../lib/withAuth"; import { getAgentSponsorStatus } from "../services/agent-sponsor"; -import { getRateLimiter, getAutumnRateLimiter } from "../services/rate-limiter"; +import { + getRateLimiter, + getAutumnRateLimiter, + getRateLimitOverride, +} from "../services/rate-limiter"; import { KEYLESS_FREE_TIER_LIMIT_MESSAGE, consumeKeylessRequest, @@ -26,7 +30,11 @@ import { AuthCreditUsageChunkRow, } from "../db/rpc"; import { AuthResponse, RateLimiterMode } from "../types"; -import { AuthCreditUsageChunk, AuthCreditUsageChunkFromTeam } from "./v1/types"; +import { + AuthCreditUsageChunk, + AuthCreditUsageChunkFromTeam, + TeamFlags, +} from "./v1/types"; import { FIRECRAWL_REST_RESOURCE, OAuthIntrospectionUnavailableError, @@ -570,14 +578,23 @@ export async function authenticateUser( * Builds the rate limiter for an authenticated team from its Autumn rate-limit * multiplier. Shared by the OAuth and API-key paths so their limiter setup * can't diverge. + * + * The org flags carry the optional per-endpoint override, so they are passed + * on to getAutumnRateLimiter, which stays the only place deciding the final + * limit. An override makes the multiplier irrelevant, so we skip fetching it + * from Autumn in that case rather than paying for a value that is discarded. */ async function buildAuthenticatedRateLimiter( teamId: string, orgId: string | null | undefined, mode: RateLimiterMode, + flags: TeamFlags, ): Promise { - const multiplier = await autumnService.getRateLimitMultiplier(teamId, orgId); - return getAutumnRateLimiter(mode, multiplier); + const multiplier = + getRateLimitOverride(mode, flags?.rateLimitOverrides) !== undefined + ? 1 + : await autumnService.getRateLimitMultiplier(teamId, orgId); + return getAutumnRateLimiter(mode, multiplier, flags); } async function supaAuthenticateUser( @@ -663,6 +680,7 @@ async function supaAuthenticateUser( teamId, chunk.org_id, mode, + chunk.flags, ); } else if (token.startsWith("fco_")) { // OAuth access token — resolve via introspection endpoint @@ -731,6 +749,7 @@ async function supaAuthenticateUser( teamId, chunk.org_id, mode, + chunk.flags, ); } else { normalizedApi = parseApi(token); @@ -761,6 +780,7 @@ async function supaAuthenticateUser( teamId, chunk.org_id, mode, + chunk.flags, ); } diff --git a/apps/api/src/controllers/v1/types.ts b/apps/api/src/controllers/v1/types.ts index 64f8a8e3d2..026ef952db 100644 --- a/apps/api/src/controllers/v1/types.ts +++ b/apps/api/src/controllers/v1/types.ts @@ -23,11 +23,13 @@ import { ProductProfile } from "../../types/product"; import { MenuProfile } from "../../types/menu"; import { threatProtectionOverrideSchema } from "../../lib/threat-protection/config"; import { auditMetadataSchema } from "../../lib/siem-logging/types"; +import type { RateLimiterMode } from "../../types"; type Format = | "markdown" | "html" | "rawHtml" + | "rawBase64" | "links" | "screenshot" | "screenshot@fullPage" @@ -435,6 +437,7 @@ const baseScrapeOptions = z.strictObject({ "markdown", "html", "rawHtml", + "rawBase64", "links", "screenshot", "screenshot@fullPage", @@ -456,6 +459,10 @@ const baseScrapeOptions = z.strictObject({ .refine( x => !x.includes("changeTracking") || x.includes("markdown"), "The changeTracking format requires the markdown format to be specified as well", + ) + .refine( + x => !x.includes("rawBase64") || x.length === 1, + "The rawBase64 format cannot be combined with other formats", ), headers: z.record(z.string(), z.string()).optional(), includeTags: z @@ -1013,6 +1020,7 @@ export type Document = { blocks?: PdfPageBlocks[]; html?: string; rawHtml?: string; + rawBase64?: string; links?: string[]; images?: string[]; screenshot?: string; @@ -1336,6 +1344,15 @@ export type TeamFlags = { >; // routes the team's new queue work to the FoundationDB backend nuqFdb?: boolean; + /** + * Per-endpoint rate-limit overrides, in requests per minute. A value here + * replaces the computed limit for that mode, so the Autumn multiplier is + * not applied. The map is sparse: a mode that is absent keeps the normal + * computation. Only a finite integer above zero is used; any other value is + * ignored. Read by getAutumnRateLimiter, so it never affects the preview + * token. + */ + rateLimitOverrides?: Partial>; } | null; export type AuthCreditUsageChunkFromTeam = Omit< diff --git a/apps/api/src/controllers/v2/types.ts b/apps/api/src/controllers/v2/types.ts index 496d00384f..ce7c0279af 100644 --- a/apps/api/src/controllers/v2/types.ts +++ b/apps/api/src/controllers/v2/types.ts @@ -458,6 +458,7 @@ export type FormatObject = | { type: "markdown" } | { type: "html" } | { type: "rawHtml" } + | { type: "rawBase64" } | { type: "links" } | { type: "images" } | { type: "summary" } @@ -702,6 +703,7 @@ const baseScrapeOptions = z.strictObject({ z.strictObject({ type: z.literal("markdown") }), z.strictObject({ type: z.literal("html") }), z.strictObject({ type: z.literal("rawHtml") }), + z.strictObject({ type: z.literal("rawBase64") }), z.strictObject({ type: z.literal("links") }), z.strictObject({ type: z.literal("images") }), z.strictObject({ type: z.literal("summary") }), @@ -737,7 +739,11 @@ const baseScrapeOptions = z.strictObject({ const hasJson = x.some(f => f.type === "json"); const hasDeterministicJson = x.some(f => f.type === "deterministicJson"); return !(hasJson && hasDeterministicJson); - }, "Cannot specify both json and deterministicJson formats"), + }, "Cannot specify both json and deterministicJson formats") + .refine( + x => !x.some(f => f.type === "rawBase64") || x.length === 1, + "The rawBase64 format cannot be combined with other formats", + ), headers: z.record(z.string(), z.string()).optional(), includeTags: z .string() @@ -1105,6 +1111,10 @@ const parseRequestSchemaBase = baseScrapeOptions.extend({ }); export const parseRequestSchema = strictWithMessage(parseRequestSchemaBase) + .refine( + x => !x.formats.some(format => format.type === "rawBase64"), + "The rawBase64 format is not supported for parse uploads", + ) .refine(waitForRefine, waitForRefineOpts) .transform(x => { const { file, ...scrapeLike } = x; @@ -1304,6 +1314,7 @@ export type Document = { blocks?: PdfPageBlocks[]; html?: string; rawHtml?: string; + rawBase64?: string; links?: string[]; images?: string[]; screenshot?: string; diff --git a/apps/api/src/scraper/scrapeURL/engines/fire-engine/index.ts b/apps/api/src/scraper/scrapeURL/engines/fire-engine/index.ts index b2a5e86ac5..06a4d3151e 100644 --- a/apps/api/src/scraper/scrapeURL/engines/fire-engine/index.ts +++ b/apps/api/src/scraper/scrapeURL/engines/fire-engine/index.ts @@ -195,24 +195,29 @@ async function performFireEngineScrape< status = scrape as FireEngineCheckStatusSuccess; } - await specialtyScrapeCheck( - logger.child({ - method: "performFireEngineScrape/specialtyScrapeCheck", - }), - status.responseHeaders, - status, - ); + const wantsRawBase64 = + hasFormatOfType(meta.options.formats, "rawBase64") !== undefined; + + if (!wantsRawBase64) { + await specialtyScrapeCheck( + logger.child({ + method: "performFireEngineScrape/specialtyScrapeCheck", + }), + status.responseHeaders, + status, + ); + } const contentType = (Object.entries(status.responseHeaders ?? {}).find( x => x[0].toLowerCase() === "content-type", ) ?? [])[1] ?? ""; - if (contentType.includes("application/json")) { + if (!wantsRawBase64 && contentType.includes("application/json")) { status.content = await getInnerJson(status.content); } - if (status.file) { + if (status.file && !wantsRawBase64) { const content = status.file.content; delete status.file; let buffer = Buffer.from(content, "base64"); @@ -318,16 +323,34 @@ export async function scrapeURLWithFireEngineChromeCDP( "engine.url": meta.url, "engine.team_id": meta.internalOptions.teamId, }); + const wantsRawBase64 = + hasFormatOfType(meta.options.formats, "rawBase64") !== undefined; + if ( + wantsRawBase64 && + ((meta.options.waitFor ?? 0) > 0 || + (meta.options.actions?.length ?? 0) > 0) + ) { + meta.logger.warn( + "rawBase64 returns the original response body; waitFor and actions are ignored.", + { + waitFor: meta.options.waitFor, + actionTypes: meta.options.actions?.map(action => action.type), + }, + ); + } const hasBranding = hasFormatOfType(meta.options.formats, "branding"); const hasAudio = hasFormatOfType(meta.options.formats, "audio"); const hasVideo = hasFormatOfType(meta.options.formats, "video"); - const shouldRunYoutubePostprocessor = youtubePostprocessor.shouldRun( - meta, - new URL(meta.rewrittenUrl ?? meta.url), - ); + const shouldRunYoutubePostprocessor = + !wantsRawBase64 && + youtubePostprocessor.shouldRun( + meta, + new URL(meta.rewrittenUrl ?? meta.url), + ); const defaultWait = hasBranding ? BRANDING_DEFAULT_WAIT_MS : 0; - const effectiveWait = - meta.options.waitFor != null && meta.options.waitFor !== 0 + const effectiveWait = wantsRawBase64 + ? 0 + : meta.options.waitFor != null && meta.options.waitFor !== 0 ? meta.options.waitFor : defaultWait; @@ -344,7 +367,7 @@ export async function scrapeURLWithFireEngineChromeCDP( : []), // Include specified actions - ...(meta.options.actions ?? []).map(action => { + ...(!wantsRawBase64 ? (meta.options.actions ?? []) : []).map(action => { const { metadata: _, ...rest } = action as InternalAction; return rest; }), @@ -407,6 +430,7 @@ export async function scrapeURLWithFireEngineChromeCDP( url: meta.rewrittenUrl ?? meta.url, scrapeId: meta.id, engine: "chrome-cdp", + ...(wantsRawBase64 ? { format: "rawBase64" as const } : {}), instantReturn: false, skipTlsVerification: meta.options.skipTlsVerification, headers: meta.options.headers, @@ -421,6 +445,7 @@ export async function scrapeURLWithFireEngineChromeCDP( timeout: meta.abort.scrapeTimeout() ?? 300000, disableSmartWaitCache: meta.internalOptions.disableSmartWaitCache, mobileProxy: meta.featureFlags.has("stealthProxy"), + autoProxy: meta.options.proxy === "auto", maxAge: meta.options.maxAge, saveScrapeResultToGCS: !meta.internalOptions.zeroDataRetention && @@ -507,6 +532,7 @@ export async function scrapeURLWithFireEngineChromeCDP( url: response.url ?? meta.url, html: response.content, + rawBase64: response.file?.content, markdown: contentType?.includes("text/markdown") ? response.content : undefined, @@ -563,6 +589,7 @@ export async function scrapeURLWithFireEngineTLSClient( geolocation: meta.options.location, disableJsDom: meta.internalOptions.v0DisableJsDom, mobileProxy: meta.featureFlags.has("stealthProxy"), + autoProxy: meta.options.proxy === "auto", timeout: meta.abort.scrapeTimeout() ?? 300000, maxAge: meta.options.maxAge, diff --git a/apps/api/src/scraper/scrapeURL/engines/fire-engine/scrape.ts b/apps/api/src/scraper/scrapeURL/engines/fire-engine/scrape.ts index cde9e009a7..03f4817025 100644 --- a/apps/api/src/scraper/scrapeURL/engines/fire-engine/scrape.ts +++ b/apps/api/src/scraper/scrapeURL/engines/fire-engine/scrape.ts @@ -30,6 +30,7 @@ const browserCookieSchema = z export type FireEngineScrapeRequestCommon = { url: string; scrapeId?: string; + format?: "html" | "rawBase64"; headers?: { [K: string]: string }; @@ -49,6 +50,7 @@ export type FireEngineScrapeRequestCommon = { geolocation?: { country?: string; languages?: string[] }; mobileProxy?: boolean; // leave it undefined if user doesn't specify + autoProxy?: boolean; timeout?: number; maxAge?: number; diff --git a/apps/api/src/scraper/scrapeURL/engines/index.ts b/apps/api/src/scraper/scrapeURL/engines/index.ts index 0401806337..8907129624 100644 --- a/apps/api/src/scraper/scrapeURL/engines/index.ts +++ b/apps/api/src/scraper/scrapeURL/engines/index.ts @@ -34,7 +34,11 @@ import { } from "../../../controllers/v2/types"; import type { PdfMetadata, PdfPageBlocks } from "./pdf/types"; import { BrandingProfile } from "../../../types/branding"; -import { AgentIndexOnlyError, BrandingNotSupportedError } from "../error"; +import { + AgentIndexOnlyError, + BrandingNotSupportedError, + NoCachedDataError, +} from "../error"; import { isUrlBlocked } from "../../WebScraper/utils/blocklist"; import { canUseExchangeForRequest, @@ -141,6 +145,7 @@ export type EngineScrapeResult = { url: string; html: string; + rawBase64?: string; markdown?: string; pages?: Array<{ pageNumber: number; markdown: string }>; blocks?: PdfPageBlocks[]; @@ -587,6 +592,27 @@ export async function buildFallbackList(meta: Meta): Promise< unsupportedFeatures: Set; }[] > { + if (hasFormatOfType(meta.options.formats, "rawBase64")) { + if (meta.internalOptions.agentIndexOnly) { + throw new AgentIndexOnlyError(); + } + + if (meta.options.minAge !== undefined) { + throw new NoCachedDataError(); + } + + if (meta.options.lockdown || (!useFireEngine && meta.mock === null)) { + return []; + } + + return [ + { + engine: "fire-engine;chrome-cdp", + unsupportedFeatures: new Set(), + }, + ]; + } + if ( !meta.internalOptions.agentIndexOnly && meta.internalOptions.forceEngine === undefined diff --git a/apps/api/src/scraper/scrapeURL/index.ts b/apps/api/src/scraper/scrapeURL/index.ts index 26e8b3cc3a..c2792ab2cc 100644 --- a/apps/api/src/scraper/scrapeURL/index.ts +++ b/apps/api/src/scraper/scrapeURL/index.ts @@ -535,6 +535,7 @@ async function scrapeURLLoopIter( const hasQuestion = hasFormatOfType(meta.options.formats, "question"); const hasHighlights = hasFormatOfType(meta.options.formats, "highlights"); const hasQuery = hasFormatOfType(meta.options.formats, "query"); + const hasRawBase64 = hasFormatOfType(meta.options.formats, "rawBase64"); const needsMarkdown = hasMarkdown || hasChangeTracking || @@ -548,7 +549,9 @@ async function scrapeURLLoopIter( const htmlSize = engineResult.html?.length ?? 0; const shouldSkipMarkdownCheck = htmlSize > MAX_HTML_SIZE_FOR_MARKDOWN_CHECK; - if ( + if (hasRawBase64) { + checkMarkdown = engineResult.rawBase64 !== undefined ? "rawBase64" : ""; + } else if ( meta.internalOptions.teamId === "sitemap" || meta.internalOptions.teamId === "robots-txt" ) { @@ -597,6 +600,9 @@ async function scrapeURLLoopIter( (engineResult.statusCode >= 200 && engineResult.statusCode < 300) || engineResult.statusCode === 304; const hasNoPageError = engineResult.error === undefined; + const hasRequiredOutput = hasRawBase64 + ? engineResult.rawBase64 !== undefined + : isLongEnough || !isGoodStatusCode; const isLikelyProxyError = [401, 403, 429].includes( engineResult.statusCode, ); @@ -622,7 +628,7 @@ async function scrapeURLLoopIter( // NOTE: TODO: what to do when status code is bad is tough... // we cannot just rely on text because error messages can be brief and not hit the limit // should we just use all the fallbacks and pick the one with the longest text? - mogery - if (isLongEnough || !isGoodStatusCode) { + if (hasRequiredOutput) { meta.logger.info("Scrape via " + engine + " deemed successful.", { factors: { isLongEnough, isGoodStatusCode, hasNoPageError }, }); @@ -954,6 +960,7 @@ async function scrapeURLLoop(meta: Meta): Promise { for (const postprocessor of postprocessors) { if ( + !hasFormatOfType(meta.options.formats, "rawBase64") && postprocessor.shouldRun( meta, new URL(engineResult.url), @@ -987,6 +994,7 @@ async function scrapeURLLoop(meta: Meta): Promise { pages: engineResult.pages, blocks: engineResult.blocks, rawHtml: engineResult.html, + rawBase64: engineResult.rawBase64, json: engineResult.json, screenshot: engineResult.screenshot, actions: engineResult.actions, diff --git a/apps/api/src/scraper/scrapeURL/scrapeURL.test.ts b/apps/api/src/scraper/scrapeURL/scrapeURL.test.ts index c5d4a6d9ab..4abfae155a 100644 --- a/apps/api/src/scraper/scrapeURL/scrapeURL.test.ts +++ b/apps/api/src/scraper/scrapeURL/scrapeURL.test.ts @@ -8,6 +8,7 @@ import { scrapeURL } from "."; import { scrapeOptions } from "../../controllers/v2/types"; import { Engine } from "./engines"; import { CostTracking } from "../../lib/cost-tracking"; +import { AgentIndexOnlyError, NoCachedDataError } from "./error"; // Mock parseMarkdown but delegate to real implementation for other tests vi.mock("../../lib/html-to-markdown", async importOriginal => { @@ -34,6 +35,92 @@ const testEnginesScreenshot: (Engine | undefined)[] = [ ]; describe("Standalone scrapeURL tests", () => { + it.each([ + ["https://example.com/raw", "text/html; charset=utf-8"], + ["https://example.com/raw.pdf", "application/pdf"], + [ + "https://example.com/raw.docx", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ], + ])( + "returns rawBase64 through fire-engine Chrome for %s", + async (url, contentType) => { + const out = await scrapeURL( + "test:raw-base64", + url, + scrapeOptions.parse({ + formats: ["rawBase64"], + useMock: "raw-base64", + fastMode: true, + blockAds: false, + proxy: "stealth", + }), + { forceEngine: "fetch", teamId: "test", orgId: null }, + new CostTracking(), + ); + + expect(out.success).toBe(true); + if (out.success) { + expect(out.document.rawBase64).toBe("PGh0bWw+cmF3PC9odG1sPg=="); + expect(out.document).not.toHaveProperty("markdown"); + expect(out.document).not.toHaveProperty("html"); + expect(out.document).not.toHaveProperty("rawHtml"); + expect(out.document.metadata.contentType).toBe(contentType); + } + }, + ); + + it("rejects a rawBase64 engine response without file content", async () => { + const out = await scrapeURL( + "test:raw-base64-missing", + "https://example.com/raw-error", + scrapeOptions.parse({ + formats: ["rawBase64"], + useMock: "raw-base64", + }), + { teamId: "test", orgId: null }, + new CostTracking(), + ); + + expect(out.success).toBe(false); + }); + + it("returns the canonical agent-index-only error for rawBase64", async () => { + const out = await scrapeURL( + "test:raw-base64-agent-index-only", + "https://example.com/raw", + scrapeOptions.parse({ + formats: ["rawBase64"], + useMock: "raw-base64", + }), + { teamId: "test", orgId: null, agentIndexOnly: true }, + new CostTracking(), + ); + + expect(out.success).toBe(false); + if (!out.success) { + expect(out.error).toBeInstanceOf(AgentIndexOnlyError); + } + }); + + it("rejects rawBase64 when minAge is set", async () => { + const out = await scrapeURL( + "test:raw-base64-min-age", + "https://example.com/raw", + scrapeOptions.parse({ + formats: ["rawBase64"], + minAge: 0, + }), + { teamId: "test", orgId: null }, + new CostTracking(), + ); + + expect(out.success).toBe(false); + if (!out.success) { + expect(out.error).toBeInstanceOf(NoCachedDataError); + } + }); + describe.each(testEngines)("Engine %s", (forceEngine: Engine | undefined) => { it("Basic scrape", async () => { const out = await scrapeURL( diff --git a/apps/api/src/scraper/scrapeURL/transformers/index.ts b/apps/api/src/scraper/scrapeURL/transformers/index.ts index 2df9b5da80..3695ae78ff 100644 --- a/apps/api/src/scraper/scrapeURL/transformers/index.ts +++ b/apps/api/src/scraper/scrapeURL/transformers/index.ts @@ -366,6 +366,7 @@ async function performLLMExtractUnlessNativeJson( function coerceFieldsToFormats(meta: Meta, document: Document): Document { const hasMarkdown = hasFormatOfType(meta.options.formats, "markdown"); const hasRawHtml = hasFormatOfType(meta.options.formats, "rawHtml"); + const hasRawBase64 = hasFormatOfType(meta.options.formats, "rawBase64"); const hasHtml = hasFormatOfType(meta.options.formats, "html"); const hasLinks = hasFormatOfType(meta.options.formats, "links"); const hasImages = hasFormatOfType(meta.options.formats, "images"); @@ -407,6 +408,14 @@ function coerceFieldsToFormats(meta: Meta, document: Document): Document { ); } + if (!hasRawBase64 && document.rawBase64 !== undefined) { + delete document.rawBase64; + } else if (hasRawBase64 && document.rawBase64 === undefined) { + meta.logger.warn( + "Request had format: rawBase64, but there was no rawBase64 field in the result.", + ); + } + if (!hasHtml && document.html !== undefined) { delete document.html; } else if (hasHtml && document.html === undefined) { @@ -661,6 +670,10 @@ export async function executeTransformers( meta: Meta, document: Document, ): Promise { + if (hasFormatOfType(meta.options.formats, "rawBase64")) { + return coerceFieldsToFormats(meta, document); + } + const executions: [string, number][] = []; for (const transformer of transformerStack) { diff --git a/apps/api/src/services/rate-limiter.test.ts b/apps/api/src/services/rate-limiter.test.ts index 709b329045..c27beda36d 100644 --- a/apps/api/src/services/rate-limiter.test.ts +++ b/apps/api/src/services/rate-limiter.test.ts @@ -1,4 +1,18 @@ +import { vi } from "vitest"; import { config } from "../config"; +import { RateLimiterMode } from "../types"; +import type { TeamFlags } from "../controllers/v1/types"; + +// The module builds a Redis client at import time. Stub ioredis so the limiter +// can be constructed without a live server. +vi.mock("ioredis", () => ({ + default: class { + defineCommand() {} + }, +})); + +import { getAutumnRateLimiter } from "./rate-limiter"; + // import { // getRateLimiter, // serverRateLimiter, @@ -369,3 +383,79 @@ import { config } from "../config"; // }); // }); // TODO: FIX + +describe("getAutumnRateLimiter", () => { + const flagsWith = (overrides: unknown): TeamFlags => + ({ rateLimitOverrides: overrides }) as TeamFlags; + + it("multiplies the base limit when there is no override", () => { + expect(getAutumnRateLimiter(RateLimiterMode.Scrape, 5).points).toBe(50); + expect(getAutumnRateLimiter(RateLimiterMode.Scrape, 5, null).points).toBe( + 50, + ); + expect( + getAutumnRateLimiter(RateLimiterMode.Scrape, 5, {} as TeamFlags).points, + ).toBe(50); + }); + + it("replaces the computed limit with the override for that mode", () => { + const limiter = getAutumnRateLimiter( + RateLimiterMode.Scrape, + 5, + flagsWith({ [RateLimiterMode.Scrape]: 42 }), + ); + + expect(limiter.points).toBe(42); + }); + + it("keeps the normal limit when only another mode is overridden", () => { + const limiter = getAutumnRateLimiter( + RateLimiterMode.Scrape, + 5, + flagsWith({ [RateLimiterMode.Map]: 42 }), + ); + + expect(limiter.points).toBe(50); + }); + + it("applies the override to a mode that has no base limit", () => { + // Research is not multiplier-scaled, so it normally uses the fallback 100. + expect(getAutumnRateLimiter(RateLimiterMode.Research, 5).points).toBe(100); + expect( + getAutumnRateLimiter( + RateLimiterMode.Research, + 5, + flagsWith({ [RateLimiterMode.Research]: 7 }), + ).points, + ).toBe(7); + }); + + it.each([ + ["zero", 0], + ["a negative number", -10], + ["a non-integer", 12.5], + ["NaN", Number.NaN], + ["Infinity", Number.POSITIVE_INFINITY], + ["a numeric string", "42"], + ["null", null], + ["an object", { rpm: 42 }], + ])("ignores an override of %s", (_label, value) => { + const limiter = getAutumnRateLimiter( + RateLimiterMode.Scrape, + 5, + flagsWith({ [RateLimiterMode.Scrape]: value }), + ); + + expect(limiter.points).toBe(50); + }); + + it("ignores an override map that is not an object", () => { + expect( + getAutumnRateLimiter(RateLimiterMode.Scrape, 5, flagsWith(null)).points, + ).toBe(50); + expect( + getAutumnRateLimiter(RateLimiterMode.Scrape, 5, flagsWith("scrape=42")) + .points, + ).toBe(50); + }); +}); diff --git a/apps/api/src/services/rate-limiter.ts b/apps/api/src/services/rate-limiter.ts index 51806d573b..7abc684847 100644 --- a/apps/api/src/services/rate-limiter.ts +++ b/apps/api/src/services/rate-limiter.ts @@ -1,6 +1,7 @@ import { RateLimiterRedis } from "rate-limiter-flexible"; import { config } from "../config"; import { RateLimiterMode } from "../types"; +import type { TeamFlags } from "../controllers/v1/types"; import Redis from "ioredis"; export const redisRateLimitClient = new Redis(config.REDIS_RATE_LIMIT_URL!, { @@ -72,16 +73,45 @@ export function getRateLimiter(mode: RateLimiterMode): RateLimiterRedis { return createRateLimiter(`${mode}`, rateLimit); } +/** + * Reads the per-minute override for one mode from the org flags. Returns + * undefined when there is no usable override. + * + * Bad config is skipped, not thrown on: the value is validated on write, but a + * broken entry must never break authentication. Only a finite integer above + * zero counts. + */ +export function getRateLimitOverride( + mode: RateLimiterMode, + overrides: unknown, +): number | undefined { + if (typeof overrides !== "object" || overrides === null) return undefined; + const value = (overrides as Record)[mode]; + if (typeof value !== "number") return undefined; + if (!Number.isInteger(value) || value <= 0) return undefined; + return value; +} + /** * Builds the per-minute rate limiter for an authenticated team from its Autumn * rate-limit multiplier: the effective limit is `base × multiplier` for * multiplier-scaled modes (default ×1). Modes without a base fall back to the * static table. + * + * An org-level override for the mode replaces that whole computation, so the + * multiplier and the base table are both discarded. Modes without an override + * keep the normal result. */ export function getAutumnRateLimiter( mode: RateLimiterMode, multiplier: number = 1, + flags?: TeamFlags, ): RateLimiterRedis { + const override = getRateLimitOverride(mode, flags?.rateLimitOverrides); + if (override !== undefined) { + return createRateLimiter(`${mode}`, override); + } + const base = BASE_RATE_LIMITS[mode]; let rateLimit: number; if (base !== undefined) { diff --git a/apps/api/v1-openapi.json b/apps/api/v1-openapi.json index 9fec50fcf6..d18d708260 100644 --- a/apps/api/v1-openapi.json +++ b/apps/api/v1-openapi.json @@ -2097,6 +2097,7 @@ "markdown", "html", "rawHtml", + "rawBase64", "links", "screenshot", "screenshot@fullPage", @@ -2105,7 +2106,7 @@ "branding" ] }, - "description": "Formats to include in the output.", + "description": "Formats to include in the output. `rawBase64` must be requested by itself.", "default": ["markdown"] }, "onlyMainContent": { @@ -2426,6 +2427,11 @@ "nullable": true, "description": "Raw HTML content of the page if `rawHtml` is in `formats`" }, + "rawBase64": { + "type": "string", + "nullable": true, + "description": "Base64-encoded original response body if `rawBase64` is in `formats`" + }, "screenshot": { "type": "string", "nullable": true, @@ -2672,6 +2678,11 @@ "nullable": true, "description": "Raw HTML content of the page if `includeRawHtml` is true" }, + "rawBase64": { + "type": "string", + "nullable": true, + "description": "Base64-encoded original response body if `rawBase64` is in `formats`" + }, "links": { "type": "array", "items": { @@ -2809,6 +2820,11 @@ "nullable": true, "description": "Raw HTML content of the page if `includeRawHtml` is true" }, + "rawBase64": { + "type": "string", + "nullable": true, + "description": "Base64-encoded original response body if `rawBase64` is in `formats`" + }, "links": { "type": "array", "items": {