fix(api): forward all highlight shadow requests (firecrawl#4036)

Author: mogery

apps/api/src/config.ts

@@ -164,7 +164,6 @@ const configSchema = z.object({
   HIGHLIGHT_MODEL_URL: z.string().optional(),
   HIGHLIGHT_MODEL_TOKEN: z.string().optional(),
   HIGHLIGHT_SHADOW_RATE: z.coerce.number().min(0).max(1).default(0),
-  HIGHLIGHT_SHADOW_MAX_INFLIGHT: z.coerce.number().int().positive().default(8),
 
   // Fire Engine
   FIRE_ENGINE_BETA_URL: z.string().optional(),

apps/api/src/search/highlight-model.ts

@@ -160,6 +160,7 @@ interface HighlightBatchOptions {
   logPayload?: boolean;
   allowLegacyFallback?: boolean;
   requestId?: string;
+  timeoutMs?: number | null;
   onFailure?: (reason: HighlightFailureReason) => void;
 }
 
@@ -175,7 +176,12 @@ async function generateHighlightsBatchRequest(
   }
 
   const controller = new AbortController();
-  const timer = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
+  const timeoutMs =
+    opts.timeoutMs === undefined ? REQUEST_TIMEOUT_MS : opts.timeoutMs;
+  const timer =
+    timeoutMs === null
+      ? undefined
+      : setTimeout(() => controller.abort(), timeoutMs);
   const start = Date.now();
   try {
     const baseUrl = config.HIGHLIGHT_MODEL_URL!.replace(/\/$/, "");
@@ -260,7 +266,9 @@ async function generateHighlightsBatchRequest(
     });
     return null;
   } finally {
-    clearTimeout(timer);
+    if (timer !== undefined) {
+      clearTimeout(timer);
+    }
   }
 }
 

apps/api/src/search/highlights-shadow.test.ts

@@ -1,7 +1,6 @@
 vi.mock("../config", () => ({
   config: {
     HIGHLIGHT_SHADOW_RATE: 1,
-    HIGHLIGHT_SHADOW_MAX_INFLIGHT: 1,
   },
 }));
 
@@ -29,7 +28,6 @@ let runSearchHighlightsShadow = createSearchHighlightsShadowRunner(logger);
 afterEach(() => {
   vi.clearAllMocks();
   config.HIGHLIGHT_SHADOW_RATE = 1;
-  config.HIGHLIGHT_SHADOW_MAX_INFLIGHT = 1;
   vi.mocked(highlightsEnvReady).mockReturnValue(true);
   runSearchHighlightsShadow = createSearchHighlightsShadowRunner(logger);
 });
@@ -77,18 +75,13 @@ describe("runSearchHighlightsShadow", () => {
     expect(fields).not.toHaveProperty("highlights");
   });
 
-  it("drops excess work instead of creating a backlog", async () => {
-    let finish!: (value: {
-      attempted: number;
-      indexHits: number;
-      replaced: number;
-      succeeded: boolean;
-    }) => void;
-    vi.mocked(runIndexedSearchHighlightsShadow).mockReturnValue(
-      new Promise(resolve => {
-        finish = resolve;
-      }),
-    );
+  it("forwards concurrent shadow work without admission drops", async () => {
+    vi.mocked(runIndexedSearchHighlightsShadow).mockResolvedValue({
+      attempted: 1,
+      indexHits: 1,
+      replaced: 1,
+      succeeded: true,
+    });
 
     const input = {
       response: {} as any,
@@ -101,18 +94,9 @@ describe("runSearchHighlightsShadow", () => {
     ).toBe("started");
     expect(
       runSearchHighlightsShadow({ ...input, requestId: "request-2" }),
-    ).toBe("dropped");
-    expect(logger.info).toHaveBeenCalledWith(
-      "Search highlights shadow dropped",
-      expect.objectContaining({
-        canonicalLog: "search/highlights-shadow",
-        outcome: "dropped",
-        reason: "max_inflight",
-      }),
-    );
-
-    finish({ attempted: 1, indexHits: 1, replaced: 1, succeeded: true });
+    ).toBe("started");
     await new Promise(resolve => setImmediate(resolve));
+    expect(runIndexedSearchHighlightsShadow).toHaveBeenCalledTimes(2);
   });
 
   it("emits the content-free failure category", async () => {

apps/api/src/search/highlights-shadow.ts

@@ -31,7 +31,7 @@ export function createSearchHighlightsShadowRunner(
   requestId: string;
   teamId: string;
   zeroDataRetention: boolean;
-}) => "skipped" | "dropped" | "started" {
+}) => "skipped" | "started" {
   let inFlight = 0;
 
   return options => {
@@ -43,20 +43,6 @@ export function createSearchHighlightsShadowRunner(
       return "skipped";
     }
 
-    if (inFlight >= config.HIGHLIGHT_SHADOW_MAX_INFLIGHT) {
-      canonicalLogger.info("Search highlights shadow dropped", {
-        canonicalLog: CANONICAL_LOG,
-        outcome: "dropped",
-        reason: "max_inflight",
-        requestId: options.requestId,
-        teamId: options.teamId,
-        inFlight,
-        maxInFlight: config.HIGHLIGHT_SHADOW_MAX_INFLIGHT,
-        sampleRate: config.HIGHLIGHT_SHADOW_RATE,
-      });
-      return "dropped";
-    }
-
     const urls = searchHighlightURLs(options.response);
     const { query, requestId, teamId } = options;
     inFlight++;
@@ -74,7 +60,6 @@ export function createSearchHighlightsShadowRunner(
           failureReason: result.failureReason,
           timeTakenMs: Date.now() - startedAt,
           inFlight,
-          maxInFlight: config.HIGHLIGHT_SHADOW_MAX_INFLIGHT,
           sampleRate: config.HIGHLIGHT_SHADOW_RATE,
         });
       })
@@ -87,7 +72,6 @@ export function createSearchHighlightsShadowRunner(
           errorType: error instanceof Error ? error.name : "unknown",
           timeTakenMs: Date.now() - startedAt,
           inFlight,
-          maxInFlight: config.HIGHLIGHT_SHADOW_MAX_INFLIGHT,
           sampleRate: config.HIGHLIGHT_SHADOW_RATE,
         });
       })

apps/api/src/search/highlights.test.ts

@@ -101,6 +101,7 @@ describe("runIndexedSearchHighlightsShadow", () => {
         logger,
         logPayload: false,
         requestId: "request-1",
+        timeoutMs: null,
         onFailure: expect.any(Function),
       },
     );

apps/api/src/search/highlights.ts

@@ -186,6 +186,7 @@ export async function runIndexedSearchHighlightsShadow(
     logger,
     logPayload: false,
     requestId,
+    timeoutMs: null,
     onFailure: reason => {
       failureReason = reason;
     },