1 files changed, 315 insertions, 0 deletions
diff --git a/castle-search/src/index.ts b/castle-search/src/index.ts
new file mode 100644
index 0000000..29cdaaf
--- /dev/null
+++ b/castle-search/src/index.ts
@@ -0,0 +1,315 @@
+/**
+ * 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");
+
+/**
+ * Say something from a command, in whichever mode pi is running.
+ *
+ * `ctx.ui.notify` is a no-op when `hasUI` is false, which means `pi -p
+ * "/search-engine"` would change the stored engine and print absolutely nothing
+ * — the command would look broken. Print mode falls back to stdout. JSON mode
+ * deliberately does not: stdout there is a structured event stream, and a stray
+ * line would corrupt it.
+ */
+const say = (ctx: { hasUI: boolean; mode: string; ui: { notify: (m: string, t?: "info" | "warning" | "error") => void } },
+  message: string,
+  level: "info" | "warning" | "error" = "info",
+): void => {
+  if (ctx.hasUI) {
+    ctx.ui.notify(message, level);
+    return;
+  }
+  if (ctx.mode === "print") console.log(message);
+};
+
+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 === "") {
+        say(
+          ctx,
+          `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.`,
+        );
+        return;
+      }
+
+      if (raw.toLowerCase() === "default" || raw.toLowerCase() === "clear") {
+        await clearStoredEngine();
+        say(ctx, `castle_cdp_search engine reset to the built-in default (${DEFAULT_ENGINE}).`);
+        return;
+      }
+
+      const id = parseEngineId(raw);
+      if (!id) {
+        say(ctx, `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();
+      say(
+        ctx,
+        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;
+      }
+    },
+  });
+}