feat(search): add staged highlights rollout (firecrawl#4041)

Author: mogery

apps/api/src/config.ts

@@ -159,11 +159,13 @@ const configSchema = z.object({
   CLICKHOUSE_ANALYTICS_URL: z.string().optional(),
   CLICKHOUSE_ANALYTICS_DATABASE: z.string().optional(),
 
-  // Search highlights (beta): highlighter service base URL. TOKEN is optional
+  // Search highlights: highlighter service base URL. TOKEN is optional
   // bearer auth for legacy/external services; the in-cluster service omits it.
   HIGHLIGHT_MODEL_URL: z.string().optional(),
   HIGHLIGHT_MODEL_TOKEN: z.string().optional(),
-  HIGHLIGHT_SHADOW_RATE: z.coerce.number().min(0).max(1).default(0),
+  // Stable percentage of non-MCP/CLI cohorts whose generated highlights are
+  // returned. The remaining eligible traffic still runs in shadow mode.
+  HIGHLIGHT_ROLLOUT_PERCENT: z.coerce.number().min(0).max(100).default(0),
 
   // Exchange (routed data sources service)
   FIRE_EXCHANGE_URL: z.url().optional(),

apps/api/src/controllers/v1/search.ts

@@ -236,6 +236,7 @@ export async function searchController(
       {
         teamId: req.auth.team_id,
         origin: req.body.origin,
+        integration: req.body.integration,
         apiKeyId: req.acuc?.api_key_id ?? null,
         flags: req.acuc?.flags ?? null,
         requestId: jobId,

apps/api/src/controllers/v1/types.ts

@@ -1319,7 +1319,6 @@ export type TeamFlags = {
   // POST /v2/search/:jobId/feedback returns 403 TEAM_OPTED_OUT when true.
   searchFeedbackOptOut?: boolean;
   researchBeta?: boolean;
-  highlightsBeta?: boolean;
   enrichBeta?: boolean;
   professionalProfileCompanyDataBeta?: boolean;
   organizationDataSourceAccess?: Record<

apps/api/src/controllers/v2/search-schema.test.ts

@@ -0,0 +1,28 @@
+import { describe, expect, it } from "vitest";
+import { searchRequestSchema } from "./types";
+
+describe("searchRequestSchema highlights", () => {
+  it("preserves an omitted value for integration and rollout selection", () => {
+    const request = searchRequestSchema.parse({ query: "firecrawl" });
+
+    expect(request.highlights).toBeUndefined();
+  });
+
+  it("allows highlights to be enabled explicitly", () => {
+    const request = searchRequestSchema.parse({
+      query: "firecrawl",
+      highlights: true,
+    });
+
+    expect(request.highlights).toBe(true);
+  });
+
+  it("allows highlights to be disabled explicitly", () => {
+    const request = searchRequestSchema.parse({
+      query: "firecrawl",
+      highlights: false,
+    });
+
+    expect(request.highlights).toBe(false);
+  });
+});

apps/api/src/controllers/v2/search.ts

@@ -219,6 +219,7 @@ export async function searchController(
       {
         teamId: req.auth.team_id,
         origin: req.body.origin,
+        integration: req.body.integration,
         apiKeyId: req.acuc?.api_key_id ?? null,
         flags: req.acuc?.flags ?? null,
         requestId: agentRequestId ?? jobId,

apps/api/src/controllers/v2/types.ts

@@ -1560,7 +1560,6 @@ export type TeamFlags = {
   debugBranding?: boolean;
   maxBrowserSessions?: number;
   researchBeta?: boolean;
-  highlightsBeta?: boolean;
   menuBeta?: boolean;
   enrichBeta?: boolean;
   professionalProfileCompanyDataBeta?: boolean;
@@ -1960,10 +1959,10 @@ export const searchRequestSchema = z
     timeout: z.int().positive().finite().prefault(60000),
     ignoreInvalidURLs: z.boolean().optional().prefault(false),
     asyncScraping: z.boolean().optional().prefault(false),
-    // Experimental: replace each result's snippet with query-relevant
-    // highlights pulled from our index (last 30 days), out-of-line from
-    // scrapeURL. Falls back to the provider snippet when the URL isn't indexed.
-    highlights: z.boolean().optional().prefault(false),
+    // Replace each result's snippet with query-relevant highlights pulled from
+    // our index. When omitted, the caller integration and rollout cohort decide
+    // whether generated highlights are returned or only run in shadow mode.
+    highlights: z.boolean().optional(),
     __searchPreviewToken: z.string().optional(),
     threatProtection: threatProtectionOverrideSchema.optional(),
     scrapeOptions: baseScrapeOptions

apps/api/src/scraper/scrapeURL/transformers/menu.ts

@@ -12,8 +12,8 @@ export async function fetchMenu(
   }
 
   // Beta gate: during beta the menu format is limited to teams with the menuBeta
-  // flag. Fail closed and silent (mirrors highlightsBeta) -- a team without the
-  // flag gets no menu field, no error, and no engine call.
+  // flag. Fail closed and silent: a team without the flag gets no menu field,
+  // no error, and no engine call.
   if (meta.internalOptions?.teamFlags?.menuBeta !== true) {
     meta.logger.info("menu format requested without menuBeta flag; skipping");
     return document;

apps/api/src/search/execute.ts

@@ -13,13 +13,18 @@ import {
   mergeScrapedContent,
   calculateScrapeCredits,
 } from "./scrape";
-import { applyIndexedSearchHighlights, highlightsEnvReady } from "./highlights";
-import { runSearchHighlightsShadow } from "./highlights-shadow";
+import {
+  highlightsEnvReady,
+  runIndexedSearchHighlights,
+  searchHighlightsMode,
+} from "./highlights";
 import { trackSearchResults, trackSearchRequest } from "../lib/tracking";
 import type { BillingMetadata } from "../services/billing/types";
 import type { ThreatProtectionPolicy } from "../lib/threat-protection/types";
 import { checkUrlsAgainstThreatPolicy } from "../lib/threat-protection/request";
 import { calculateThreatScanCredits } from "../lib/scrape-billing";
+import { config } from "../config";
+import { logger as rootLogger } from "../lib/logger";
 
 interface SearchOptions {
   query: string;
@@ -42,6 +47,7 @@ interface SearchOptions {
 interface SearchContext {
   teamId: string;
   origin: string;
+  integration?: string | null;
   apiKeyId: number | null;
   flags: TeamFlags;
   requestId: string;
@@ -241,32 +247,47 @@ export async function executeSearch(
     }
   }
 
-  // Experimental highlights beta: replace provider snippets with index-backed
-  // highlights. Gated on (1) the request opting in, (2) the team's highlightsBeta
-  // flag, and (3) all required envs being present (index DB, GCS, model). Any
-  // gate failing => silently keep the provider snippets.
+  // Every eligible request runs the same index-backed highlight path. MCP/CLI,
+  // explicit opt-ins, and rollout-selected cohorts await and apply the result;
+  // everyone else runs it in shadow without delaying or mutating the response.
   // Runs after scraping (mergeScrapedContent rebuilds the result objects, so
   // highlight mutations must come last to survive). Uses the user's original
   // query, not the domain-filtered upstream query.
-  const shouldApplyHighlights =
-    options.highlights &&
-    flags?.highlightsBeta === true &&
-    highlightsEnvReady();
-  if (shouldApplyHighlights) {
-    await applyIndexedSearchHighlights(
+  if (zeroDataRetention !== true && !isZDR && highlightsEnvReady()) {
+    const mode = searchHighlightsMode({
+      requested: options.highlights,
+      origin: context.origin,
+      integration: context.integration,
+      cohortKey:
+        context.apiKeyId !== null
+          ? `api-key:${context.apiKeyId}`
+          : `team:${context.teamId}`,
+      rolloutPercent: config.HIGHLIGHT_ROLLOUT_PERCENT,
+    });
+    const highlightRun = runIndexedSearchHighlights(
       searchResponse,
       query,
       logger,
-      context.requestId,
+      {
+        mode,
+        requestId: context.requestId,
+        teamId: context.teamId,
+      },
     );
-  } else {
-    runSearchHighlightsShadow({
-      response: searchResponse,
-      query,
-      requestId: context.requestId,
-      teamId,
-      zeroDataRetention: zeroDataRetention === true || isZDR === true,
-    });
+
+    if (mode === "apply") {
+      await highlightRun;
+    } else {
+      void highlightRun.catch(error => {
+        rootLogger.warn("Search highlights shadow failed", {
+          canonicalLog: "search/highlights",
+          mode,
+          requestId: context.requestId,
+          teamId: context.teamId,
+          errorType: error instanceof Error ? error.name : "unknown",
+        });
+      });
+    }
   }
 
   const scrapeFormats = scrapeOptions?.formats