path: root/src/search.ts

/**
 * Search orchestration: resolve a browser, borrow a page, navigate, extract,
 * hand the page back.
 *
 * The connection is process-wide and lazy. It is *not* opened by the extension
 * factory — pi runs factories in invocations that never start a session, and
 * the docs are explicit that sockets must not start there. It is opened on the
 * first search and closed by an idempotent `session_shutdown`.
 */

import { BrowserConnection, PageSession, raceAbort } from "./cdp.ts";
import {
  type CdpTarget,
  describeCdpTarget,
  describeResolutionFailure,
  resolveCdpTarget,
} from "./endpoint.ts";
import { describeEngineFailure, describeEngineSource, resolveEngine } from "./engine-target.ts";
import { readStoredEngine } from "./engine-store.ts";
import {
  ENGINE_IDS as allEngineIdList,
  type EngineDefinition,
  type EngineId,
  type Recency,
  describeRecencySupport,
  enginesSupportingRecency,
  getEngine,
} from "./engines.ts";

const allEngineIds = (): readonly EngineId[] => allEngineIdList;
import {
  BrowserUnavailableError,
  NoResultsError,
  RecencyUnsupportedError,
  SearchChallengeError,
} from "./errors.ts";
import { buildProbeScript, type PageProbe, type SearchResult } from "./extract.ts";
import { readStoredTarget } from "./target-store.ts";

/** Per-search wall clock, covering connect, navigate and extract together. */
export const SEARCH_TIMEOUT_MS = 20_000;
/** How many searches may drive the shared browser at once. */
export const MAX_CONCURRENT_SEARCHES = 2;
/** Results sometimes land a beat after the load event; re-read a couple of times. */
const EXTRACT_RETRIES = 2;
const EXTRACT_RETRY_DELAY_MS = 700;

export type { Recency } from "./engines.ts";

export type SearchRequest = {
  query: string;
  numResults: number;
  recency?: Recency | undefined;
  /** Per-call engine override; falls back to the configured default. */
  engine?: string | null | undefined;
};

export type SearchOutcome = {
  results: SearchResult[];
  searchUrl: string;
  /** The URL the page settled on, which differs from searchUrl after a redirect. */
  finalUrl: string;
  endpoint: string;
  engine: EngineDefinition;
  /** Where the engine choice came from, for the "via" line in the output. */
  engineSource: string;
};

/** Build the SERP URL for a request. Exported for tests. */
export const buildSearchUrl = (engine: EngineDefinition, request: SearchRequest): string =>
  engine.buildUrl(request.query, request.numResults, request.recency);

/**
 * Resolve the engine, refusing rather than silently ignoring a recency window
 * the chosen engine cannot express.
 *
 * Dropping an unsupported filter would be the same failure mode as Google's
 * `tbs=qdr:` — the caller asks for recent results, gets whatever the engine
 * ranked, and has no way to tell. Naming the engines that *do* support the
 * window makes the error actionable, since the model can retry with one.
 */
export const selectEngine = async (
  request: SearchRequest,
): Promise<{ engine: EngineDefinition; source: string }> => {
  const resolution = resolveEngine(request.engine, process.env, await readStoredEngine());
  if (resolution.kind === "invalid") {
    throw new BrowserUnavailableError(describeEngineFailure(resolution));
  }
  const engine = getEngine(resolution.id);

  if (request.recency && engine.recency[request.recency] === undefined) {
    const alternatives = enginesSupportingRecency(request.recency);
    throw new RecencyUnsupportedError(
      `${engine.label} cannot restrict results to the past ${request.recency}. ` +
        `It supports: ${describeRecencySupport(engine)}. ` +
        (alternatives.length > 0
          ? `Engines that support "${request.recency}": ${alternatives.join(", ")}. ` +
            `Retry with engine set to one of those, or drop the recency parameter.`
          : `No configured engine supports that window; drop the recency parameter.`),
    );
  }

  return { engine, source: describeEngineSource(resolution.source) };
};

/** For status reporting — resolves the configured engine without searching. */
export const describeConfiguredEngine = async (): Promise<{ id: EngineId; source: string; label: string }> => {
  const resolution = resolveEngine(null, process.env, await readStoredEngine());
  if (resolution.kind === "invalid") {
    return { id: "google", source: describeEngineFailure(resolution), label: "unresolved" };
  }
  return {
    id: resolution.id,
    source: describeEngineSource(resolution.source),
    label: getEngine(resolution.id).label,
  };
};

// ── connection lifecycle ───────────────────────────────────────────────────

let connection: BrowserConnection | null = null;
/** In-flight connect, so two concurrent searches share one dial rather than racing. */
let connecting: Promise<BrowserConnection> | null = null;

const resolveTarget = async (): Promise<CdpTarget> => {
  const resolution = resolveCdpTarget(process.env, await readStoredTarget());
  if (resolution.kind === "invalid") {
    throw new BrowserUnavailableError(describeResolutionFailure(resolution));
  }
  return resolution.target;
};

/**
 * Return a live connection, dialling one if there isn't one or the last one
 * died. Reconnecting transparently matters here: castle recycles Chrome hourly,
 * so a session that searched an hour ago is holding a dead socket.
 */
const ensureConnection = async (signal?: AbortSignal): Promise<BrowserConnection> => {
  if (connection && connection.isOpen) return connection;
  if (connection) {
    connection.close();
    connection = null;
  }

  if (!connecting) {
    // Assigned synchronously, before the first await inside: two searches
    // starting in the same tick must share one dial, not open two sockets and
    // leak whichever loses the assignment race.
    connecting = (async () => {
      const target = await resolveTarget();
      const conn = await BrowserConnection.connect(target);
      connection = conn;
      return conn;
    })().finally(() => {
      connecting = null;
    });
  }

  // The caller's signal is raced here rather than passed into connect(). A
  // shared dial must not be cancelled by whichever caller happened to start it
  // — one search pressing Esc would otherwise fail the other search waiting on
  // the same connection. connect() is bounded by its own timeouts regardless.
  return raceAbort(connecting, signal);
};

/** Idempotent: safe to call from `session_shutdown` however many times it fires. */
export const disconnect = (): void => {
  connection?.close();
  connection = null;
};

/** For status reporting — never dials. */
export const describeConfiguredEndpoint = async (): Promise<string> => {
  const resolution = resolveCdpTarget(process.env, await readStoredTarget());
  return resolution.kind === "invalid"
    ? describeResolutionFailure(resolution)
    : describeCdpTarget(resolution.target);
};

// ── concurrency cap ────────────────────────────────────────────────────────

let active = 0;
const waiting: Array<() => void> = [];

/** Give a held slot to the next waiter, or return it to the pool. */
const handOff = (): void => {
  const next = waiting.shift();
  // Hand the slot straight over rather than decrementing and letting the waiter
  // re-check, so a queued search cannot be overtaken by a newly arriving one.
  if (next) next();
  else active--;
};

const acquire = async (signal?: AbortSignal): Promise<() => void> => {
  if (active < MAX_CONCURRENT_SEARCHES) {
    active++;
  } else {
    let resolver!: () => void;
    const queued = new Promise<void>((resolve) => {
      resolver = resolve;
    });
    waiting.push(resolver);
    try {
      await raceAbort(queued, signal);
      // The slot was handed over by handOff(), which already counted it.
    } catch (e) {
      // An abort while queued must not strand the slot. Either we are still in
      // the queue — drop out of it — or the slot was handed to us in the same
      // tick the abort landed, in which case give it back. Without this, two
      // cancelled-while-queued searches would wedge the semaphore permanently.
      const index = waiting.indexOf(resolver);
      if (index !== -1) waiting.splice(index, 1);
      else handOff();
      throw e;
    }
  }
  let released = false;
  return () => {
    if (released) return;
    released = true;
    handOff();
  };
};

// ── the search itself ──────────────────────────────────────────────────────

export type ProgressReporter = (phase: SearchPhase, note: string) => void;
export type SearchPhase = "queued" | "connecting" | "opening" | "navigating" | "extracting";

const sleep = (ms: number, signal?: AbortSignal): Promise<void> =>
  raceAbort(new Promise<void>((resolve) => setTimeout(resolve, ms)), signal);

/**
 * The other engines, named in failure messages. A blocked engine is the case
 * where the agent most needs to know it has somewhere else to go, and telling
 * it inside the error beats hoping it remembers the guidelines.
 */
const otherEngines = (current: EngineId): readonly string[] =>
  allEngineIds().filter((id) => id !== current);

/**
 * Run one search. Throws on every failure — see errors.ts for why a returned
 * error object would be worse than useless here.
 */
export const search = async (
  request: SearchRequest,
  signal: AbortSignal,
  report: ProgressReporter,
): Promise<SearchOutcome> => {
  // Reported before acquiring, so a search that dies waiting for a slot is
  // distinguishable from one that died reaching the browser. The deadline is
  // the caller's and covers the queue wait too — with a cap of two and searches
  // that take a couple of seconds, waiting is the rare case, and a tool call
  // that silently takes twice its stated deadline would be worse.
  // Resolved before acquiring a slot: an unknown engine name or an unsupported
  // recency window is a caller error, and should not wait behind other searches
  // or touch the browser at all.
  const { engine, source: engineSource } = await selectEngine(request);
  const probeScript = buildProbeScript(engine.extraction);

  report("queued", "waiting for a search slot");
  const release = await acquire(signal);
  try {
    report("connecting", "resolving browser");
    const conn = await ensureConnection(signal);
    const endpoint = describeCdpTarget(conn.target);

    report("opening", `opening a page on ${endpoint}`);
    const page = await PageSession.open(conn, signal);
    const searchUrl = buildSearchUrl(engine, request);

    try {
      report("navigating", `searching ${engine.label} for "${request.query}"`);
      await page.navigate(searchUrl, signal);

      report("extracting", "reading results");
      let probe = await page.evaluate<PageProbe>(probeScript, signal);
      for (let attempt = 0; attempt < EXTRACT_RETRIES; attempt++) {
        if (probe?.challenge || (probe?.results?.length ?? 0) > 0) break;
        await sleep(EXTRACT_RETRY_DELAY_MS, signal);
        probe = await page.evaluate<PageProbe>(probeScript, signal);
      }

      if (!probe) {
        throw new BrowserUnavailableError("the page returned nothing — the extraction script did not run");
      }

      if (probe.challenge) {
        throw new SearchChallengeError({
          challenge: probe.challenge.kind,
          // The page the human must clear is where the browser *ended up*, which
          // after a consent redirect is not the URL we asked for.
          pageUrl: probe.url || searchUrl,
          endpoint,
          engine: engine.label,
          alternatives: otherEngines(engine.id),
          detail: probe.challenge.detail || undefined,
        });
      }

      if (probe.results.length === 0) {
        // "container missing entirely" and "container present but empty" are
        // different faults, and saying which saves the next person a browser
        // session.
        const diagnosis =
          probe.container === "found"
            ? `${engine.label}'s results container was rendered but nothing could be parsed out of it, which ` +
              `suggests its markup changed and castle_cdp_search's extractor for ${engine.id} needs updating.`
            : `${engine.label} served a page with no results area at all. This is often a soft rate-limit: as a ` +
              `block decays an engine stops serving its challenge page and returns an empty result page instead, ` +
              `which is indistinguishable from a genuine zero-hit search (observed on Google lasting ~90 minutes ` +
              `after heavy querying). It can also mean the query truly has no hits.`;
        throw new NoResultsError(
          `No results could be read from ${probe.url || searchUrl} (page title: "${probe.title}"). ${diagnosis} ` +
            `Before concluding the extractor is broken, retry with a different engine (${otherEngines(engine.id).join(", ")}); ` +
            `if those also come back empty the query is the problem, and if only this one does, it is rate limiting — ` +
            `wait rather than retrying in a loop.`,
        );
      }

      return {
        results: probe.results.slice(0, request.numResults),
        searchUrl,
        finalUrl: probe.url || searchUrl,
        endpoint,
        engine,
        engineSource,
      };
    } finally {
      // Always, including the challenge path. Leaving the tab open would let the
      // operator solve the captcha in place, but it would also litter a browser
      // someone else is using with abandoned tabs on every failure; the error
      // carries the URL instead, and the cookie it sets is profile-wide.
      await page.close();
    }
  } finally {
    release();
  }
};