1 files changed, 198 insertions, 0 deletions
diff --git a/src/index.ts b/src/index.ts
new file mode 100644
index 0000000..ad0ba6f
--- /dev/null
+++ b/src/index.ts
@@ -0,0 +1,198 @@
+/**
+ * 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 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 { messageOf } from "./errors.ts";
+import type { SearchResult } from "./extract.ts";
+import { formatResults, truncationNote } from "./format.ts";
+import {
+  MAX_CONCURRENT_SEARCHES,
+  SEARCH_TIMEOUT_MS,
+  type Recency,
+  type SearchPhase,
+  describeConfiguredEndpoint,
+  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(["day", "week", "month", "year"] as const, {
+      description: "Restrict results to pages published within this window. Omit for no time limit.",
+    }),
+  ),
+});
+
+/** Exported so `isToolCallEventType<"castle_cdp_search", CastleCdpSearchInput>` can type it. */
+export type CastleCdpSearchInput = {
+  query: string;
+  numResults?: number;
+  recency?: Recency;
+};
+
+export type CastleCdpSearchDetails = {
+  query: string;
+  numResults: number;
+  recency: Recency | 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): CastleCdpSearchDetails => ({
+  query,
+  numResults,
+  recency,
+  phase: "queued",
+  note: "starting",
+  endpoint: null,
+  searchUrl: null,
+  results: [],
+  truncated: false,
+});
+
+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();
+  });
+
+  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. 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); 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.",
+      "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.",
+      "If castle_cdp_search fails with a SearchChallengeError, the browser is sitting behind a captcha or " +
+        "consent wall that only a person can clear: stop searching, tell the user the URL from the error, and " +
+        "ask them to complete the challenge in that browser before you retry.",
+    ],
+    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 details = initialDetails(query, numResults, recency);
+
+      // 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 }, combined, report);
+
+        const body = formatResults({
+          query,
+          searchUrl: outcome.searchUrl,
+          finalUrl: outcome.finalUrl,
+          endpoint: outcome.endpoint,
+          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`;
+        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.
+        // The endpoint is worth naming: "which browser failed" is most of the
+        // diagnosis when the same config can point at three different machines.
+        const endpoint = await describeConfiguredEndpoint().catch(() => "unknown endpoint");
+        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}" — browser ${endpoint}, ` +
+          `up to ${MAX_CONCURRENT_SEARCHES} concurrent searches, ${SEARCH_TIMEOUT_MS / 1000}s deadline covering the whole call]`;
+        throw e2;
+      }
+    },
+  });
+}