fix(api): move highlight shadow preprocessing off API (firecrawl#4035)

mogery · web-flow · commit bb9175fdf11f · 2026-07-15T11:54:22.000+02:00
* fix(api): move highlight shadow preprocessing off api

* test(api): separate indexed shadow suite
diff --git a/apps/api/src/search/highlight-model.test.ts b/apps/api/src/search/highlight-model.test.ts
@@ -5,7 +5,10 @@ vi.mock("../config", () => ({
   },
 }));
 
-import { generateHighlightsBatch } from "./highlight-model";
+import {
+  generateHighlightsBatch,
+  generateHighlightsIndexedBatch,
+} from "./highlight-model";
 import { config } from "../config";
 
 const logger = {
@@ -33,6 +36,35 @@ afterEach(() => {
 });
 
 describe("generateHighlightsBatch", () => {
+  it("posts lightweight index references to the indexed Stage 1 endpoint", async () => {
+    const fetchMock = mockFetchOnce({ pages: [] });
+
+    await generateHighlightsIndexedBatch(
+      "q1",
+      [
+        {
+          id: "0",
+          url: "https://first.test/path",
+          indexObject: "index-object.json",
+        },
+      ],
+      { logger, logPayload: false, requestId: "request-1" },
+    );
+
+    const [url, init] = fetchMock.mock.calls[0];
+    expect(url).toBe("https://highlight.test/batch_highlight_indexed");
+    expect(JSON.parse(init.body)).toEqual({
+      query: "q1",
+      pages: [
+        {
+          id: "0",
+          url: "https://first.test/path",
+          indexObject: "index-object.json",
+        },
+      ],
+    });
+  });
+
   it("posts every page to one /batch_highlight call with the bearer token", async () => {
     const fetchMock = mockFetchOnce({ pages: [] });
 
diff --git a/apps/api/src/search/highlight-model.ts b/apps/api/src/search/highlight-model.ts
@@ -61,6 +61,12 @@ interface HighlightPage {
   markdown: string;
 }
 
+export interface HighlightIndexedPage {
+  id: string;
+  url: string;
+  indexObject: string;
+}
+
 // Result for one page in the batch: the service's highlight entries plus the
 // reassembled markdown document.
 interface HighlightResult {
@@ -149,16 +155,20 @@ function requestHeaders(requestId?: string): Record<string, string> {
  * its provider snippet without discarding successful pages. Returns null when
  * the whole call fails.
  */
-export async function generateHighlightsBatch(
+interface HighlightBatchOptions {
+  logger: Logger;
+  logPayload?: boolean;
+  allowLegacyFallback?: boolean;
+  requestId?: string;
+  onFailure?: (reason: HighlightFailureReason) => void;
+}
+
+async function generateHighlightsBatchRequest(
+  endpoint: "/batch_highlight" | "/batch_highlight_indexed",
   query: string,
-  pages: HighlightPage[],
-  opts: {
-    logger: Logger;
-    logPayload?: boolean;
-    allowLegacyFallback?: boolean;
-    requestId?: string;
-    onFailure?: (reason: HighlightFailureReason) => void;
-  },
+  pages: Array<HighlightPage | HighlightIndexedPage>,
+  opts: HighlightBatchOptions,
+  legacyPages?: HighlightPage[],
 ): Promise<Map<string, HighlightResult> | null> {
   if (pages.length === 0) {
     return new Map();
@@ -170,7 +180,7 @@ export async function generateHighlightsBatch(
   try {
     const baseUrl = config.HIGHLIGHT_MODEL_URL!.replace(/\/$/, "");
     const res = await fetchBatchWithRetry(
-      `${baseUrl}/batch_highlight`,
+      `${baseUrl}${endpoint}`,
       {
         method: "POST",
         headers: requestHeaders(opts.requestId),
@@ -186,6 +196,7 @@ export async function generateHighlightsBatch(
       // endpoint, whose /batch_highlight contract uses {requests: [...]}. Keep
       // highlights available until infra switches the URL to GCP Stage 1.
       if (
+        legacyPages &&
         opts.allowLegacyFallback !== false &&
         (res.status === 400 || res.status === 404)
       ) {
@@ -196,7 +207,7 @@ export async function generateHighlightsBatch(
         return await generateLegacyHighlightsBatch(
           baseUrl,
           query,
-          pages,
+          legacyPages,
           controller.signal,
           opts,
         );
@@ -253,6 +264,33 @@ export async function generateHighlightsBatch(
   }
 }
 
+export async function generateHighlightsBatch(
+  query: string,
+  pages: HighlightPage[],
+  opts: HighlightBatchOptions,
+): Promise<Map<string, HighlightResult> | null> {
+  return generateHighlightsBatchRequest(
+    "/batch_highlight",
+    query,
+    pages,
+    opts,
+    pages,
+  );
+}
+
+export async function generateHighlightsIndexedBatch(
+  query: string,
+  pages: HighlightIndexedPage[],
+  opts: Omit<HighlightBatchOptions, "allowLegacyFallback">,
+): Promise<Map<string, HighlightResult> | null> {
+  return generateHighlightsBatchRequest(
+    "/batch_highlight_indexed",
+    query,
+    pages,
+    { ...opts, allowLegacyFallback: false },
+  );
+}
+
 async function generateLegacyHighlightsBatch(
   baseUrl: string,
   query: string,
diff --git a/apps/api/src/search/highlights-shadow.test.ts b/apps/api/src/search/highlights-shadow.test.ts
@@ -7,11 +7,16 @@ vi.mock("../config", () => ({
 
 vi.mock("./highlights", () => ({
   highlightsEnvReady: vi.fn(() => true),
-  applySearchHighlights: vi.fn(),
+  searchHighlightURLs: vi.fn(() => []),
+  runIndexedSearchHighlightsShadow: vi.fn(),
 }));
 
 import { config } from "../config";
-import { applySearchHighlights, highlightsEnvReady } from "./highlights";
+import {
+  highlightsEnvReady,
+  runIndexedSearchHighlightsShadow,
+  searchHighlightURLs,
+} from "./highlights";
 import { createSearchHighlightsShadowRunner } from "./highlights-shadow";
 
 const logger = {
@@ -31,7 +36,7 @@ afterEach(() => {
 
 describe("runSearchHighlightsShadow", () => {
   it("runs without applying results and emits a content-free canonical log", async () => {
-    vi.mocked(applySearchHighlights).mockResolvedValue({
+    vi.mocked(runIndexedSearchHighlightsShadow).mockResolvedValue({
       attempted: 10,
       indexHits: 7,
       replaced: 6,
@@ -49,17 +54,12 @@ describe("runSearchHighlightsShadow", () => {
     ).toBe("started");
     await new Promise(resolve => setImmediate(resolve));
 
-    expect(applySearchHighlights).toHaveBeenCalledWith(
-      {},
+    expect(searchHighlightURLs).toHaveBeenCalledWith({});
+    expect(runIndexedSearchHighlightsShadow).toHaveBeenCalledWith(
+      [],
       "private query",
       expect.objectContaining({ silent: true }),
-      {
-        applyResults: false,
-        suppressSummaryLog: true,
-        suppressPayloadLog: true,
-        allowLegacyFallback: false,
-        requestId: "request-1",
-      },
+      "request-1",
     );
     expect(logger.info).toHaveBeenCalledWith(
       "Search highlights shadow completed",
@@ -84,7 +84,7 @@ describe("runSearchHighlightsShadow", () => {
       replaced: number;
       succeeded: boolean;
     }) => void;
-    vi.mocked(applySearchHighlights).mockReturnValue(
+    vi.mocked(runIndexedSearchHighlightsShadow).mockReturnValue(
       new Promise(resolve => {
         finish = resolve;
       }),
@@ -116,7 +116,7 @@ describe("runSearchHighlightsShadow", () => {
   });
 
   it("emits the content-free failure category", async () => {
-    vi.mocked(applySearchHighlights).mockResolvedValue({
+    vi.mocked(runIndexedSearchHighlightsShadow).mockResolvedValue({
       attempted: 3,
       indexHits: 1,
       replaced: 0,
@@ -165,6 +165,6 @@ describe("runSearchHighlightsShadow", () => {
       runSearchHighlightsShadow({ ...input, zeroDataRetention: false }),
     ).toBe("skipped");
 
-    expect(applySearchHighlights).not.toHaveBeenCalled();
+    expect(runIndexedSearchHighlightsShadow).not.toHaveBeenCalled();
   });
 });
diff --git a/apps/api/src/search/highlights-shadow.ts b/apps/api/src/search/highlights-shadow.ts
@@ -2,7 +2,11 @@ import { createLogger, type Logger } from "winston";
 import type { SearchV2Response } from "../lib/entities";
 import { config } from "../config";
 import { logger as rootLogger } from "../lib/logger";
-import { applySearchHighlights, highlightsEnvReady } from "./highlights";
+import {
+  highlightsEnvReady,
+  runIndexedSearchHighlightsShadow,
+  searchHighlightURLs,
+} from "./highlights";
 
 const CANONICAL_LOG = "search/highlights-shadow";
 const shadowLogger = createLogger({ silent: true });
@@ -53,21 +57,17 @@ export function createSearchHighlightsShadowRunner(
       return "dropped";
     }
 
+    const urls = searchHighlightURLs(options.response);
+    const { query, requestId, teamId } = options;
     inFlight++;
     const startedAt = Date.now();
-    void applySearchHighlights(options.response, options.query, shadowLogger, {
-      applyResults: false,
-      suppressSummaryLog: true,
-      suppressPayloadLog: true,
-      allowLegacyFallback: false,
-      requestId: options.requestId,
-    })
+    void runIndexedSearchHighlightsShadow(urls, query, shadowLogger, requestId)
       .then(result => {
         canonicalLogger.info("Search highlights shadow completed", {
           canonicalLog: CANONICAL_LOG,
           outcome: result.succeeded ? "completed" : "failed",
-          requestId: options.requestId,
-          teamId: options.teamId,
+          requestId,
+          teamId,
           attempted: result.attempted,
           indexHits: result.indexHits,
           wouldReplace: result.replaced,
@@ -82,8 +82,8 @@ export function createSearchHighlightsShadowRunner(
         canonicalLogger.warn("Search highlights shadow failed", {
           canonicalLog: CANONICAL_LOG,
           outcome: "failed",
-          requestId: options.requestId,
-          teamId: options.teamId,
+          requestId,
+          teamId,
           errorType: error instanceof Error ? error.name : "unknown",
           timeTakenMs: Date.now() - startedAt,
           inFlight,
diff --git a/apps/api/src/search/highlights.test.ts b/apps/api/src/search/highlights.test.ts
@@ -35,12 +35,21 @@ vi.mock("../scraper/scrapeURL/lib/removeUnwantedElements", () => ({
 
 vi.mock("./highlight-model", () => ({
   generateHighlightsBatch: vi.fn(),
+  generateHighlightsIndexedBatch: vi.fn(),
 }));
 
-import { generateHighlightsBatch } from "./highlight-model";
+import {
+  generateHighlightsBatch,
+  generateHighlightsIndexedBatch,
+} from "./highlight-model";
 import { config } from "../config";
 import { indexGetRecent5 } from "../db/rpc";
-import { applySearchHighlights, highlightsEnvReady } from "./highlights";
+import {
+  applySearchHighlights,
+  highlightsEnvReady,
+  runIndexedSearchHighlightsShadow,
+  searchHighlightURLs,
+} from "./highlights";
 
 const logger = {
   info: vi.fn(),
@@ -56,6 +65,54 @@ afterEach(() => {
   vi.clearAllMocks();
 });
 
+describe("runIndexedSearchHighlightsShadow", () => {
+  it("resolves lightweight index references without loading page content", async () => {
+    vi.mocked(generateHighlightsIndexedBatch).mockResolvedValue(
+      new Map([["0", { highlights: [], markdown: "shadow highlight" }]]),
+    );
+    const response = {
+      web: [{ url: "https://first.test", description: "fallback" }],
+      news: [{ url: "https://second.test", snippet: "fallback" }],
+    } as any;
+    const urls = searchHighlightURLs(response);
+
+    const result = await runIndexedSearchHighlightsShadow(
+      urls,
+      "query",
+      logger,
+      "request-1",
+    );
+
+    expect(generateHighlightsIndexedBatch).toHaveBeenCalledWith(
+      "query",
+      [
+        {
+          id: "0",
+          url: "https://first.test",
+          indexObject: "index:https://first.test.json",
+        },
+        {
+          id: "1",
+          url: "https://second.test",
+          indexObject: "index:https://second.test.json",
+        },
+      ],
+      {
+        logger,
+        logPayload: false,
+        requestId: "request-1",
+        onFailure: expect.any(Function),
+      },
+    );
+    expect(result).toEqual({
+      attempted: 2,
+      indexHits: 2,
+      replaced: 1,
+      succeeded: true,
+    });
+  });
+});
+
 describe("applySearchHighlights", () => {
   it("enables the in-cluster service without requiring a bearer token", () => {
     const token = config.HIGHLIGHT_MODEL_TOKEN;
diff --git a/apps/api/src/search/highlights.ts b/apps/api/src/search/highlights.ts
