Sybil-2/server/src/search/brave.ts

import { buildBrowserLikeRequestHeaders } from "../browser-fetch-headers.js";
import { env } from "../env.js";

const BRAVE_WEB_SEARCH_URL = "https://api.search.brave.com/res/v1/web/search";
const BRAVE_SEARCH_TIMEOUT_MS = 12_000;
const DEFAULT_BRAVE_REQUEST_INTERVAL_MS = 1_000;
const RATE_LIMIT_INTERVAL_SAFETY_RATIO = 0.05;
const MIN_RATE_LIMIT_INTERVAL_SAFETY_MS = 2;
const RATE_LIMIT_RESET_SAFETY_MS = 50;
const MAX_RATE_LIMIT_RETRIES = 3;
const MAX_RATE_LIMIT_RETRY_DELAY_MS = 8_000;

type RateLimitPolicy = {
  limit: number;
  windowSeconds: number;
};

let requestIntervalMs = addIntervalSafety(DEFAULT_BRAVE_REQUEST_INTERVAL_MS);
let lastRequestAtMs = 0;
let nextRequestAtMs = 0;
let quotaUnavailableUntilMs = 0;
let requestQueue = Promise.resolve();

export type BraveSearchOptions = {
  numResults: number;
  includeDomains?: string[];
  excludeDomains?: string[];
};

export type BraveSearchResult = {
  title: string | null;
  url: string | null;
  publishedDate: string | null;
  author: string | null;
  summary: string | null;
  text: string | null;
  highlights: string[];
};

export type BraveSearchResponse = {
  query: string;
  requestId: string | null;
  results: BraveSearchResult[];
};

function clipText(input: string, maxCharacters: number) {
  return input.length <= maxCharacters ? input : `${input.slice(0, maxCharacters)}...`;
}

function compactWhitespace(input: string) {
  return input.replace(/\r/g, "").replace(/[ \t]+\n/g, "\n").replace(/\n{3,}/g, "\n\n").replace(/\s+/g, " ").trim();
}

function requireBraveSearchApiKey() {
  if (!env.BRAVE_SEARCH_API_KEY) {
    throw new Error("BRAVE_SEARCH_API_KEY not set");
  }
  return env.BRAVE_SEARCH_API_KEY;
}

function sleep(milliseconds: number) {
  return new Promise<void>((resolve) => setTimeout(resolve, milliseconds));
}

function addIntervalSafety(intervalMs: number) {
  return intervalMs + Math.max(MIN_RATE_LIMIT_INTERVAL_SAFETY_MS, Math.ceil(intervalMs * RATE_LIMIT_INTERVAL_SAFETY_RATIO));
}

function parseCommaSeparatedNumbers(value: string | null) {
  if (!value) return [];
  return value.split(",").map((part) => Number(part.trim())).map((number) => (Number.isFinite(number) ? number : null));
}

function parseRateLimitPolicy(value: string | null): RateLimitPolicy[] {
  if (!value) return [];
  return value.split(",").flatMap((part) => {
    const match = part.trim().match(/^(\d+)\s*;\s*w=(\d+)$/i);
    if (!match) return [];
    const limit = Number(match[1]);
    const windowSeconds = Number(match[2]);
    return limit > 0 && windowSeconds > 0 ? [{ limit, windowSeconds }] : [];
  });
}

function getBurstPolicyIndex(policies: RateLimitPolicy[]) {
  if (!policies.length) return null;
  let burstIndex = 0;
  for (let index = 1; index < policies.length; index += 1) {
    if (policies[index]!.windowSeconds < policies[burstIndex]!.windowSeconds) burstIndex = index;
  }
  return burstIndex;
}

function updateRateLimitState(headers: Headers) {
  const policies = parseRateLimitPolicy(headers.get("x-ratelimit-policy"));
  const burstIndex = getBurstPolicyIndex(policies);
  if (burstIndex === null) return;

  const burstPolicy = policies[burstIndex]!;
  const learnedIntervalMs = addIntervalSafety(Math.ceil((burstPolicy.windowSeconds * 1_000) / burstPolicy.limit));
  if (learnedIntervalMs < requestIntervalMs && lastRequestAtMs > 0) {
    nextRequestAtMs = Math.min(nextRequestAtMs, lastRequestAtMs + learnedIntervalMs);
  }
  requestIntervalMs = learnedIntervalMs;

  const remaining = parseCommaSeparatedNumbers(headers.get("x-ratelimit-remaining"));
  const resetSeconds = parseCommaSeparatedNumbers(headers.get("x-ratelimit-reset"));
  for (let index = 0; index < policies.length; index += 1) {
    if ((remaining[index] ?? null) === null || remaining[index]! >= 1 || (resetSeconds[index] ?? 0) <= 0) continue;
    const unavailableUntilMs = Date.now() + resetSeconds[index]! * 1_000 + RATE_LIMIT_RESET_SAFETY_MS;
    if (index === burstIndex) {
      nextRequestAtMs = Math.max(nextRequestAtMs, unavailableUntilMs);
    } else {
      quotaUnavailableUntilMs = Math.max(quotaUnavailableUntilMs, unavailableUntilMs);
    }
  }
}

function assertLongTermQuotaAvailable() {
  if (quotaUnavailableUntilMs <= Date.now()) {
    quotaUnavailableUntilMs = 0;
    return;
  }
  const resetSeconds = Math.ceil((quotaUnavailableUntilMs - Date.now()) / 1_000);
  throw new Error(`Brave Search API long-term quota is exhausted; reset is expected in ${resetSeconds} seconds.`);
}

async function waitForRateLimitSlot() {
  const reservation = requestQueue.then(async () => {
    while (true) {
      assertLongTermQuotaAvailable();
      const waitMs = nextRequestAtMs - Date.now();
      if (waitMs <= 0) break;
      await sleep(waitMs);
    }
    lastRequestAtMs = Date.now();
    nextRequestAtMs = lastRequestAtMs + requestIntervalMs;
  });
  requestQueue = reservation.catch(() => undefined);
  await reservation;
}

function get429RetryDelayMs(headers: Headers, retryNumber: number) {
  const remaining = parseCommaSeparatedNumbers(headers.get("x-ratelimit-remaining"));
  const resetSeconds = parseCommaSeparatedNumbers(headers.get("x-ratelimit-reset"));
  const exhaustedResetSeconds = resetSeconds.filter((reset, index): reset is number => reset !== null && (remaining[index] ?? 0) < 1);
  const headerDelayMs = exhaustedResetSeconds.length ? Math.max(...exhaustedResetSeconds) * 1_000 : 0;
  const exponentialDelayMs = 2 ** retryNumber * 1_000;
  const delayMs = Math.max(headerDelayMs + RATE_LIMIT_RESET_SAFETY_MS, exponentialDelayMs);
  return delayMs <= MAX_RATE_LIMIT_RETRY_DELAY_MS ? delayMs : null;
}

async function fetchBrave(url: URL) {
  const apiKey = requireBraveSearchApiKey();
  for (let attempt = 0; attempt <= MAX_RATE_LIMIT_RETRIES; attempt += 1) {
    await waitForRateLimitSlot();

    const controller = new AbortController();
    const timeout = setTimeout(() => controller.abort(), BRAVE_SEARCH_TIMEOUT_MS);
    let response: Response;
    try {
      response = await fetch(url, {
        signal: controller.signal,
        headers: {
          ...buildBrowserLikeRequestHeaders("application/json"),
          "X-Subscription-Token": apiKey,
        },
      });
    } finally {
      clearTimeout(timeout);
    }

    updateRateLimitState(response.headers);
    if (response.status !== 429 || attempt === MAX_RATE_LIMIT_RETRIES) return response;

    const retryDelayMs = get429RetryDelayMs(response.headers, attempt);
    await response.arrayBuffer();
    if (retryDelayMs === null) {
      throw new Error("Brave Search API rate limit quota is exhausted beyond the retry window.");
    }
    await sleep(retryDelayMs);
  }

  throw new Error("Brave Search API request failed after rate-limit retries.");
}

function normalizeDomain(input: string) {
  const trimmed = input.trim().toLowerCase();
  if (!trimmed) return null;

  try {
    const parsed = new URL(trimmed.includes("://") ? trimmed : `https://${trimmed}`);
    return parsed.hostname.replace(/^www\./, "");
  } catch {
    return trimmed.split(/[/?#]/, 1)[0]?.replace(/^www\./, "") || null;
  }
}

function normalizeDomains(input: string[] | undefined) {
  return Array.from(new Set((input ?? []).map(normalizeDomain).filter((domain): domain is string => Boolean(domain))));
}

function hostnameMatchesDomain(urlRaw: string | null, domain: string) {
  if (!urlRaw) return false;
  try {
    const hostname = new URL(urlRaw).hostname.toLowerCase().replace(/^www\./, "");
    return hostname === domain || hostname.endsWith(`.${domain}`);
  } catch {
    return false;
  }
}

function filterResultsByDomains(results: BraveSearchResult[], options: BraveSearchOptions) {
  const includeDomains = normalizeDomains(options.includeDomains);
  const excludeDomains = normalizeDomains(options.excludeDomains);
  return results.filter((result) => {
    if (includeDomains.length && !includeDomains.some((domain) => hostnameMatchesDomain(result.url, domain))) return false;
    if (excludeDomains.some((domain) => hostnameMatchesDomain(result.url, domain))) return false;
    return true;
  });
}

function buildBraveQuery(query: string, options: BraveSearchOptions) {
  const includeDomains = normalizeDomains(options.includeDomains);
  const excludeDomains = normalizeDomains(options.excludeDomains);
  const includeClause =
    includeDomains.length === 0
      ? ""
      : includeDomains.length === 1
        ? `site:${includeDomains[0]}`
        : `(${includeDomains.map((domain) => `site:${domain}`).join(" OR ")})`;
  const excludeClause = excludeDomains.map((domain) => `-site:${domain}`).join(" ");
  return [query, includeClause, excludeClause].filter(Boolean).join(" ");
}

function buildSearchUrl(query: string, options: BraveSearchOptions) {
  const url = new URL(BRAVE_WEB_SEARCH_URL);
  url.searchParams.set("q", buildBraveQuery(query, options));
  url.searchParams.set("count", String(options.numResults));
  url.searchParams.set("safesearch", "moderate");
  url.searchParams.set("result_filter", "web");
  url.searchParams.set("text_decorations", "false");
  url.searchParams.set("extra_snippets", "true");
  return url;
}

function stringOrNull(value: unknown) {
  if (typeof value !== "string") return null;
  const normalized = compactWhitespace(value);
  return normalized || null;
}

function stringArray(value: unknown) {
  if (!Array.isArray(value)) return [];
  return value.filter((item): item is string => typeof item === "string").map(compactWhitespace).filter(Boolean);
}

function mapWebResult(result: any): BraveSearchResult {
  const description = stringOrNull(result?.description);
  const extraSnippets = stringArray(result?.extra_snippets);
  const snippets = [description, ...extraSnippets].filter((snippet): snippet is string => Boolean(snippet));
  const combinedText = snippets.join("\n\n");

  return {
    title: stringOrNull(result?.title),
    url: stringOrNull(result?.url),
    publishedDate: stringOrNull(result?.page_age),
    author: stringOrNull(result?.profile?.name) ?? stringOrNull(result?.article?.author),
    summary: description ? clipText(description, 1_400) : null,
    text: combinedText ? clipText(combinedText, 700) : null,
    highlights: snippets.slice(0, 3).map((snippet) => clipText(snippet, 280)),
  };
}

export async function searchBrave(query: string, options: BraveSearchOptions): Promise<BraveSearchResponse> {
  const url = buildSearchUrl(query, options);
  const response = await fetchBrave(url);

  if (!response.ok) {
    await response.arrayBuffer();
    throw new Error(`Brave Search API request failed with status ${response.status}.`);
  }

  const contentType = response.headers.get("content-type")?.toLowerCase() ?? "";
  if (!contentType.includes("application/json")) {
    await response.arrayBuffer();
    throw new Error(`Brave Search API returned ${contentType || "unknown content type"}.`);
  }

  const data: any = await response.json();
  const results = Array.isArray(data?.web?.results) ? data.web.results.map(mapWebResult) : [];
  return {
    query,
    requestId: response.headers.get("x-request-id"),
    results: filterResultsByDomains(results, options).slice(0, options.numResults),
  };
}

export function resetBraveRateLimitStateForTests() {
  requestIntervalMs = addIntervalSafety(DEFAULT_BRAVE_REQUEST_INTERVAL_MS);
  lastRequestAtMs = 0;
  nextRequestAtMs = 0;
  quotaUnavailableUntilMs = 0;
  requestQueue = Promise.resolve();
}