path: root/src/index.ts

/**
 * castle-cdp-search — web search for pi through a Chrome that is already
 * running somewhere on the wire network, driven over CDP.
 *
 * It does not launch a browser. It attaches to one that a person or a launchd
 * job started, using the same endpoint configuration as pi-browser-harness
 * (`BU_CDP_HTTP` > `/browser-target` > castle), so a single choice governs both
 * packages and they can never end up driving different browsers.
 *
 * The engine is switchable the same way: `CCS_SEARCH_ENGINE` > `/search-engine`
 * > Google, plus a per-call `engine` parameter so the agent can fall back when
 * one engine starts serving captchas.
 *
 * The tool is called `castle_cdp_search`, not `web_search`, on purpose: it is
 * meant to sit alongside pi-web-access rather than shadow it. The model picks
 * between them, so both stay available in one session.
 */

import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, formatSize, truncateHead } from "@earendil-works/pi-coding-agent";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { StringEnum } from "@earendil-works/pi-ai";
import { Type } from "typebox";

import { clearStoredEngine, storedEngineLocation, writeStoredEngine } from "./engine-store.ts";
import { DEFAULT_ENGINE, ENGINE_ENV_VAR } from "./engine-target.ts";
import {
  ENGINE_IDS,
  RECENCY_VALUES,
  type EngineId,
  type Recency,
  allEngines,
  describeRecencySupport,
  engineAliases,
  parseEngineId,
} from "./engines.ts";
import { messageOf } from "./errors.ts";
import type { SearchResult } from "./extract.ts";
import { formatResults, truncationNote } from "./format.ts";
import {
  MAX_CONCURRENT_SEARCHES,
  SEARCH_TIMEOUT_MS,
  type SearchPhase,
  describeConfiguredEndpoint,
  describeConfiguredEngine,
  disconnect,
  search,
} from "./search.ts";

const DEFAULT_NUM_RESULTS = 10;
const MIN_NUM_RESULTS = 1;
const MAX_NUM_RESULTS = 20;

const parameters = Type.Object({
  query: Type.String({
    description: "The search query, phrased as you would type it into a search engine.",
    minLength: 1,
  }),
  numResults: Type.Optional(
    Type.Integer({
      description: `How many results to return (${MIN_NUM_RESULTS}-${MAX_NUM_RESULTS}, default ${DEFAULT_NUM_RESULTS}).`,
      minimum: MIN_NUM_RESULTS,
      maximum: MAX_NUM_RESULTS,
      default: DEFAULT_NUM_RESULTS,
    }),
  ),
  // StringEnum, not Type.Union of Type.Literal: Google's API rejects the
  // anyOf/const shape typebox emits for unions.
  recency: Type.Optional(
    StringEnum(RECENCY_VALUES, {
      description:
        "Restrict results to pages published within this window. Omit for no time limit. " +
        "Not every engine supports every window — bing has no year, brave has no time filter at all.",
    }),
  ),
  engine: Type.Optional(
    StringEnum(ENGINE_IDS, {
      description:
        "Which search engine to use for this call. Omit to use the configured default. " +
        "Switch engines when one returns a challenge or no results — they have independent indexes and blocks.",
    }),
  ),
});

/** Exported so `isToolCallEventType<"castle_cdp_search", CastleCdpSearchInput>` can type it. */
export type CastleCdpSearchInput = {
  query: string;
  numResults?: number;
  recency?: Recency;
  engine?: EngineId;
};

export type CastleCdpSearchDetails = {
  query: string;
  numResults: number;
  recency: Recency | null;
  engine: EngineId | null;
  phase: SearchPhase | "done";
  note: string;
  endpoint: string | null;
  searchUrl: string | null;
  results: SearchResult[];
  truncated: boolean;
};

const initialDetails = (
  query: string,
  numResults: number,
  recency: Recency | null,
  engine: EngineId | null,
): CastleCdpSearchDetails => ({
  query,
  numResults,
  recency,
  engine,
  phase: "queued",
  note: "starting",
  endpoint: null,
  searchUrl: null,
  results: [],
  truncated: false,
});

/** The engine table, rendered for `/search-engine` and for the tool description. */
const engineTable = (): string =>
  allEngines()
    .map((e) => `  ${e.id.padEnd(11)} ${e.label.padEnd(14)} recency: ${describeRecencySupport(e).padEnd(24)} ${e.note}`)
    .join("\n");

export default function (pi: ExtensionAPI): void {
  // Nothing is dialled here. pi runs extension factories in invocations that
  // never start a session, so opening a socket from a factory would leave one
  // behind on, say, `pi --list-models`. The connection is made on first use and
  // torn down below.

  pi.on("session_shutdown", async () => {
    disconnect();
  });

  // The engine analogue of pi-browser-harness's /browser-target. Persisted
  // machine-wide for the same reason: the choice should outlive the session
  // that made it.
  pi.registerCommand("search-engine", {
    description: "Show or set the search engine castle_cdp_search uses",
    handler: async (args, ctx) => {
      const raw = args.trim();
      const current = await describeConfiguredEngine();

      if (raw === "") {
        ctx.ui.notify(
          `castle_cdp_search engine: ${current.id} (${current.source})\n\n${engineTable()}\n\n` +
            `Set with: /search-engine <name>   Clear with: /search-engine default\n` +
            `Stored in ${storedEngineLocation()}; ${ENGINE_ENV_VAR} overrides it for one process.`,
          "info",
        );
        return;
      }

      if (raw.toLowerCase() === "default" || raw.toLowerCase() === "clear") {
        await clearStoredEngine();
        ctx.ui.notify(`castle_cdp_search engine reset to the built-in default (${DEFAULT_ENGINE}).`, "info");
        return;
      }

      const id = parseEngineId(raw);
      if (!id) {
        ctx.ui.notify(`Unknown search engine "${raw}". Expected one of: ${engineAliases().join(", ")}`, "error");
        return;
      }

      await writeStoredEngine(id);
      // Saying so explicitly, because the variable silently wins otherwise and
      // the user would reasonably assume the command had taken effect.
      const pinned = process.env[ENGINE_ENV_VAR]?.trim();
      ctx.ui.notify(
        pinned
          ? `Saved ${id}, but ${ENGINE_ENV_VAR}=${pinned} is set and overrides it for this process.`
          : `castle_cdp_search will use ${id}.`,
        pinned ? "warning" : "info",
      );
    },
  });

  pi.registerTool({
    name: "castle_cdp_search",
    label: "Castle Search",
    description:
      "Search the web using a real Chrome browser running on the private network, driven over the Chrome " +
      "DevTools Protocol. Returns ranked results with titles, URLs and snippets. Because it uses a real " +
      "logged-in browser rather than a search API, it reaches pages that block datacentre traffic. It " +
      "returns search results only — it does not fetch or read the linked pages, so follow up with a " +
      "fetch/read tool for full page content. Supports four engines with independent indexes and " +
      `independent rate limits (${ENGINE_IDS.join(", ")}); pass engine to switch when one is blocked. ` +
      `Output is truncated at ${formatSize(DEFAULT_MAX_BYTES)} or ${DEFAULT_MAX_LINES} lines.`,
    promptSnippet:
      "Search the web through a real Chrome on the private network (castle_cdp_search); google/duckduckgo/bing/brave, returns titles, URLs and snippets",
    promptGuidelines: [
      "Use castle_cdp_search to find pages on the open web when you need current information, documentation, " +
        "or sources you do not already have — it drives a real signed-in Chrome, so it works on sites that " +
        "reject scripted clients.",
      "castle_cdp_search returns search results only. To read a result, follow it up with a tool that fetches " +
        "page content; do not treat the snippet as the full page.",
      "If castle_cdp_search fails with a SearchChallengeError or returns no results, retry the same query with " +
        "its engine parameter set to a different engine (google, duckduckgo, bing, brave) before giving up — " +
        "they have independent indexes and independent blocks, and a challenge on one says nothing about the others.",
      "Only ask the user for help with a castle_cdp_search captcha once more than one engine has failed; " +
        "when you do, give them the URL from the error so they can clear it in that browser.",
      "Pass recency to castle_cdp_search when the answer depends on how recent a page is, and keep numResults " +
        "small (5-10) unless you genuinely need a wide sweep. Not every engine supports every window: bing has " +
        "no year, and brave has no time filter, so castle_cdp_search will tell you to switch engines rather than " +
        "silently ignoring the request.",
    ],
    parameters,

    async execute(_toolCallId, params, signal, onUpdate, _ctx) {
      const query = params.query.trim();
      if (query === "") throw new Error("query must not be empty");

      const numResults = Math.min(
        MAX_NUM_RESULTS,
        Math.max(MIN_NUM_RESULTS, params.numResults ?? DEFAULT_NUM_RESULTS),
      );
      const recency = (params.recency ?? null) as Recency | null;
      const engine = (params.engine ?? null) as EngineId | null;

      const details = initialDetails(query, numResults, recency, engine);

      // One deadline covering connect + navigate + extract, plus the user's own
      // abort so Esc drops an in-flight navigation rather than waiting it out.
      const timeout = AbortSignal.timeout(SEARCH_TIMEOUT_MS);
      const combined = signal ? AbortSignal.any([signal, timeout]) : timeout;

      const report = (phase: SearchPhase, note: string): void => {
        details.phase = phase;
        details.note = note;
        // Keep the TUI moving during the seconds Chrome spends navigating.
        onUpdate?.({ content: [{ type: "text", text: `${note}…` }], details: { ...details } });
      };

      try {
        const outcome = await search(
          { query, numResults, recency: recency ?? undefined, engine },
          combined,
          report,
        );

        const body = formatResults({
          query,
          searchUrl: outcome.searchUrl,
          finalUrl: outcome.finalUrl,
          endpoint: outcome.endpoint,
          engine: outcome.engine.label,
          engineSource: outcome.engineSource,
          results: outcome.results,
        });

        const truncation = truncateHead(body, { maxLines: DEFAULT_MAX_LINES, maxBytes: DEFAULT_MAX_BYTES });
        const text = truncation.truncated
          ? truncation.content +
            truncationNote(
              truncation.outputLines,
              truncation.totalLines,
              formatSize(truncation.outputBytes),
              formatSize(truncation.totalBytes),
            )
          : truncation.content;

        details.phase = "done";
        details.note = `${outcome.results.length} results from ${outcome.engine.id}`;
        details.engine = outcome.engine.id;
        details.endpoint = outcome.endpoint;
        details.searchUrl = outcome.searchUrl;
        details.results = outcome.results;
        details.truncated = truncation.truncated;

        return { content: [{ type: "text", text }], details };
      } catch (e) {
        // Rethrow, always. pi only marks a tool result as an error when execute()
        // throws — returning `{ error }` would read to the model as a search that
        // simply found nothing, which is the one wrong conclusion to invite here.
        const endpoint = await describeConfiguredEndpoint().catch(() => "unknown endpoint");
        const configured = await describeConfiguredEngine().catch(() => null);
        const e2 = e instanceof Error ? e : new Error(messageOf(e));
        // Naming the phase turns "it timed out" into something actionable:
        // "queued" means the cap was the bottleneck, "navigating" means the
        // browser or the network was.
        e2.message =
          `${e2.message} [castle_cdp_search: failed during "${details.phase}" — engine ` +
          `${engine ?? configured?.id ?? "unknown"}, browser ${endpoint}, up to ${MAX_CONCURRENT_SEARCHES} ` +
          `concurrent searches, ${SEARCH_TIMEOUT_MS / 1000}s deadline covering the whole call]`;
        throw e2;
      }
    },
  });
}