fix(search): retry transient highlight failures (firecrawl#4000)

* fix(search): retry transient highlight failures

* fix(api): harden highlight retry handling

* fix(api): reject malformed highlight batches

Author: mogery

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

@@ -27,6 +27,7 @@ function mockFetchOnce(body: unknown, ok = true, status = 200) {
 }
 
 afterEach(() => {
+  vi.useRealTimers();
   vi.unstubAllGlobals();
   vi.clearAllMocks();
 });
@@ -172,17 +173,95 @@ describe("generateHighlightsBatch", () => {
 
   it("returns null when the service errors", async () => {
     mockFetchOnce({ error: "boom" }, false, 500);
+    const onFailure = vi.fn();
 
     const out = await generateHighlightsBatch(
       "q",
       [{ id: "0", markdown: "md" }],
-      { logger },
+      { logger, onFailure },
     );
 
     expect(out).toBeNull();
+    expect(fetch).toHaveBeenCalledTimes(2);
+    expect(onFailure).toHaveBeenCalledWith("http_5xx");
     expect(logger.warn).toHaveBeenCalled();
   });
 
+  it("retries one transient server failure", async () => {
+    const fetchMock = vi
+      .fn()
+      .mockResolvedValueOnce({
+        ok: false,
+        status: 503,
+        text: async () => "unavailable",
+      })
+      .mockResolvedValueOnce({
+        ok: true,
+        status: 200,
+        json: async () => ({ pages: [] }),
+        text: async () => "",
+      });
+    vi.stubGlobal("fetch", fetchMock);
+
+    const out = await generateHighlightsBatch(
+      "q",
+      [{ id: "0", markdown: "md" }],
+      { logger },
+    );
+
+    expect(fetchMock).toHaveBeenCalledTimes(2);
+    expect(out).toEqual(new Map());
+  });
+
+  it.each([null, []])(
+    "classifies malformed JSON response %j as invalid",
+    async body => {
+      mockFetchOnce(body);
+      const onFailure = vi.fn();
+
+      const out = await generateHighlightsBatch(
+        "q",
+        [{ id: "0", markdown: "md" }],
+        { logger, onFailure },
+      );
+
+      expect(out).toBeNull();
+      expect(onFailure).toHaveBeenCalledWith("invalid_response");
+    },
+  );
+
+  it("does not retry after the request deadline aborts during backoff", async () => {
+    vi.useFakeTimers();
+    const fetchMock = vi.fn().mockImplementation(
+      () =>
+        new Promise(resolve =>
+          setTimeout(
+            () =>
+              resolve({
+                ok: false,
+                status: 503,
+                text: async () => "unavailable",
+              }),
+            29_990,
+          ),
+        ),
+    );
+    vi.stubGlobal("fetch", fetchMock);
+    const onFailure = vi.fn();
+
+    const resultPromise = generateHighlightsBatch(
+      "q",
+      [{ id: "0", markdown: "md" }],
+      { logger, onFailure },
+    );
+    await vi.advanceTimersByTimeAsync(30_000);
+    const out = await resultPromise;
+
+    expect(out).toBeNull();
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+    expect(onFailure).toHaveBeenCalledWith("timeout");
+  });
+
   it("falls back to legacy per-page calls while the old service URL is configured", async () => {
     const fetchMock = vi
       .fn()

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

@@ -10,6 +10,27 @@ import { config } from "../config";
 // the in-cluster GCP Stage 1 service relies on cluster network isolation.
 
 const REQUEST_TIMEOUT_MS = 30000;
+const MAX_BATCH_ATTEMPTS = 2;
+const RETRY_DELAY_MS = 50;
+
+export type HighlightFailureReason =
+  | "timeout"
+  | "network"
+  | "http_4xx"
+  | "http_5xx"
+  | "invalid_response"
+  | "unknown";
+
+class HighlightHttpError extends Error {
+  constructor(
+    readonly status: number,
+    body: string,
+  ) {
+    super(`highlight model HTTP ${status}: ${body.slice(0, 200)}`);
+  }
+}
+
+class HighlightInvalidResponseError extends Error {}
 
 // One highlight entry as returned by the service. Field semantics are
 // intentionally left undocumented here.
@@ -47,6 +68,68 @@ interface HighlightResult {
   markdown: string;
 }
 
+function failureReason(error: unknown): HighlightFailureReason {
+  if (error instanceof HighlightHttpError) {
+    return error.status >= 500 ? "http_5xx" : "http_4xx";
+  }
+  if (error instanceof SyntaxError) {
+    return "invalid_response";
+  }
+  if (error instanceof HighlightInvalidResponseError) {
+    return "invalid_response";
+  }
+  if (error instanceof Error && error.name === "AbortError") {
+    return "timeout";
+  }
+  if (error instanceof TypeError) {
+    return "network";
+  }
+  return "unknown";
+}
+
+function waitForRetry(signal: AbortSignal): Promise<void> {
+  if (signal.aborted) {
+    return Promise.reject(signal.reason);
+  }
+
+  return new Promise((resolve, reject) => {
+    const onAbort = () => {
+      clearTimeout(timer);
+      reject(signal.reason);
+    };
+    const timer = setTimeout(() => {
+      signal.removeEventListener("abort", onAbort);
+      resolve();
+    }, RETRY_DELAY_MS);
+    signal.addEventListener("abort", onAbort, { once: true });
+  });
+}
+
+async function fetchBatchWithRetry(
+  url: string,
+  init: RequestInit,
+  signal: AbortSignal,
+): Promise<Response> {
+  let lastError: unknown;
+  for (let attempt = 1; attempt <= MAX_BATCH_ATTEMPTS; attempt++) {
+    try {
+      const response = await fetch(url, init);
+      if (response.status < 500 || attempt === MAX_BATCH_ATTEMPTS) {
+        return response;
+      }
+      await response.body?.cancel().catch(() => undefined);
+      lastError = new HighlightHttpError(response.status, "");
+    } catch (error) {
+      lastError = error;
+      if (signal.aborted || attempt === MAX_BATCH_ATTEMPTS) {
+        throw error;
+      }
+    }
+    await waitForRetry(signal);
+  }
+  throw lastError;
+}
+
 function requestHeaders(requestId?: string): Record<string, string> {
   const headers: Record<string, string> = {
     "Content-Type": "application/json",
@@ -74,6 +157,7 @@ export async function generateHighlightsBatch(
     logPayload?: boolean;
     allowLegacyFallback?: boolean;
     requestId?: string;
+    onFailure?: (reason: HighlightFailureReason) => void;
   },
 ): Promise<Map<string, HighlightResult> | null> {
   if (pages.length === 0) {
@@ -85,12 +169,16 @@ export async function generateHighlightsBatch(
   const start = Date.now();
   try {
     const baseUrl = config.HIGHLIGHT_MODEL_URL!.replace(/\/$/, "");
-    const res = await fetch(`${baseUrl}/batch_highlight`, {
-      method: "POST",
-      headers: requestHeaders(opts.requestId),
-      body: JSON.stringify({ query, pages }),
-      signal: controller.signal,
-    });
+    const res = await fetchBatchWithRetry(
+      `${baseUrl}/batch_highlight`,
+      {
+        method: "POST",
+        headers: requestHeaders(opts.requestId),
+        body: JSON.stringify({ query, pages }),
+        signal: controller.signal,
+      },
+      controller.signal,
+    );
 
     if (!res.ok) {
       const body = await res.text().catch(() => "");
@@ -113,14 +201,23 @@ export async function generateHighlightsBatch(
           opts,
         );
       }
-      throw new Error(
-        `highlight model HTTP ${res.status}: ${body.slice(0, 200)}`,
-      );
+      throw new HighlightHttpError(res.status, body);
     }
 
-    const data = (await res.json()) as HighlightBatchResponse;
+    const data: unknown = await res.json();
+    if (
+      typeof data !== "object" ||
+      data === null ||
+      Array.isArray(data) ||
+      ("pages" in data && !Array.isArray(data.pages))
+    ) {
+      throw new HighlightInvalidResponseError(
+        "highlight model returned an invalid response",
+      );
+    }
     const results = new Map<string, HighlightResult>();
-    for (const page of data.pages ?? []) {
+    for (const page of (data as HighlightBatchResponse).pages ?? []) {
+      if (typeof page !== "object" || page === null) continue;
       if (typeof page.id !== "string" || !page.output) continue;
       results.set(page.id, {
         highlights: page.output.highlights ?? [],
@@ -145,6 +242,7 @@ export async function generateHighlightsBatch(
 
     return results;
   } catch (error) {
+    opts.onFailure?.(failureReason(error));
     opts.logger.warn("query highlights batch failed", {
       canonicalLog: "search/highlights",
       error: error instanceof Error ? error.message : String(error),

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

@@ -115,6 +115,33 @@ describe("runSearchHighlightsShadow", () => {
     await new Promise(resolve => setImmediate(resolve));
   });
 
+  it("emits the content-free failure category", async () => {
+    vi.mocked(applySearchHighlights).mockResolvedValue({
+      attempted: 3,
+      indexHits: 1,
+      replaced: 0,
+      succeeded: false,
+      failureReason: "network",
+    });
+
+    runSearchHighlightsShadow({
+      response: {} as any,
+      query: "private query",
+      requestId: "request-1",
+      teamId: "team-1",
+      zeroDataRetention: false,
+    });
+    await new Promise(resolve => setImmediate(resolve));
+
+    expect(logger.info).toHaveBeenCalledWith(
+      "Search highlights shadow completed",
+      expect.objectContaining({
+        outcome: "failed",
+        failureReason: "network",
+      }),
+    );
+  });
+
   it("does not shadow ZDR, disabled, or unavailable requests", () => {
     const input = {
       response: {} as any,

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

@@ -71,6 +71,7 @@ export function createSearchHighlightsShadowRunner(
           attempted: result.attempted,
           indexHits: result.indexHits,
           wouldReplace: result.replaced,
+          failureReason: result.failureReason,
           timeTakenMs: Date.now() - startedAt,
           inFlight,
           maxInFlight: config.HIGHLIGHT_SHADOW_MAX_INFLIGHT,

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

@@ -95,7 +95,10 @@ describe("applySearchHighlights", () => {
           markdown: "markdown:<main>index:https://second.test.json</main>",
         },
       ],
-      { logger },
+      {
+        logger,
+        onFailure: expect.any(Function),
+      },
     );
     expect(response.web[0].description).toBe("first highlight");
     expect(response.news[0].snippet).toBe("second highlight");
@@ -184,6 +187,7 @@ describe("applySearchHighlights", () => {
         logger,
         logPayload: false,
         allowLegacyFallback: false,
+        onFailure: expect.any(Function),
       },
     );
     expect(logger.info).not.toHaveBeenCalledWith(

apps/api/src/search/highlights.ts

@@ -11,6 +11,7 @@ import { parseMarkdown } from "../lib/html-to-markdown";
 import { htmlTransform } from "../scraper/scrapeURL/lib/removeUnwantedElements";
 import type { ScrapeOptions } from "../controllers/v2/types";
 import { generateHighlightsBatch } from "./highlight-model";
+import type { HighlightFailureReason } from "./highlight-model";
 import { config } from "../config";
 
 // How far back into the index we're willing to reach for highlight source text.
@@ -142,6 +143,7 @@ export async function applySearchHighlights(
   indexHits: number;
   replaced: number;
   succeeded: boolean;
+  failureReason?: HighlightFailureReason;
 }> {
   const start = Date.now();
   const applyResults = options.applyResults ?? true;
@@ -194,6 +196,7 @@ export async function applySearchHighlights(
   // empty response only falls back the corresponding provider snippet.
   let replaced = 0;
   let succeeded = true;
+  let failureReason: HighlightFailureReason | undefined;
   if (indexHits > 0) {
     const results = await generateHighlightsBatch(
       query,
@@ -206,9 +209,18 @@ export async function applySearchHighlights(
             logger,
             logPayload: !options.suppressPayloadLog,
             allowLegacyFallback: options.allowLegacyFallback,
-            requestId: options.requestId,
+            ...(options.requestId ? { requestId: options.requestId } : {}),
+            onFailure: reason => {
+              failureReason = reason;
+            },
           }
-        : { logger, requestId: options.requestId },
+        : {
+            logger,
+            ...(options.requestId ? { requestId: options.requestId } : {}),
+            onFailure: reason => {
+              failureReason = reason;
+            },
+          },
     );
     succeeded = results !== null;
     if (results) {
@@ -238,5 +250,6 @@ export async function applySearchHighlights(
     indexHits,
     replaced,
     succeeded,
+    ...(failureReason ? { failureReason } : {}),
   };
 }