all repos — sz-chrome-remote-cdp @ master

pi-browser-harness, patched to drive a Chrome on another host over CDP, with castle_cdp_search web search

castle-search/src/index.ts (view raw)

 1
 2
 3
 4
 5
 6
 7
 8
 9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
 100
 101
 102
 103
 104
 105
 106
 107
 108
 109
 110
 111
 112
 113
 114
 115
 116
 117
 118
 119
 120
 121
 122
 123
 124
 125
 126
 127
 128
 129
 130
 131
 132
 133
 134
 135
 136
 137
 138
 139
 140
 141
 142
 143
 144
 145
 146
 147
 148
 149
 150
 151
 152
 153
 154
 155
 156
 157
 158
 159
 160
 161
 162
 163
 164
 165
 166
 167
 168
 169
 170
 171
 172
 173
 174
 175
 176
 177
 178
 179
 180
 181
 182
 183
 184
 185
 186
 187
 188
 189
 190
 191
 192
 193
 194
 195
 196
 197
 198
 199
 200
 201
 202
 203
 204
 205
 206
 207
 208
 209
 210
 211
 212
 213
 214
 215
 216
 217
 218
 219
 220
 221
 222
 223
 224
 225
 226
 227
 228
 229
 230
 231
 232
 233
 234
 235
 236
 237
 238
 239
 240
 241
 242
 243
 244
 245
 246
 247
 248
 249
 250
 251
 252
 253
 254
 255
 256
 257
 258
 259
 260
 261
 262
 263
 264
 265
 266
 267
 268
 269
 270
 271
 272
 273
 274
 275
 276
 277
 278
 279
 280
 281
 282
 283
 284
 285
 286
 287
 288
 289
 290
 291
 292
 293
 294
 295
 296
 297
 298
 299
 300
 301
 302
 303
 304
 305
 306
 307
 308
 309
 310
 311
 312
 313
 314
 315
 316
 317
 318
 319
 320
 321
 322
 323
 324
 325
 326
 327
 328
 329
 330
 331
 332
 333
 334
 335
 336
 337
 338
 339
 340
 341
 342
 343
 344
 345
 346
 347
 348
 349
 350
/**
 * 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,
  resolveLink,
  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_resolve",
    label: "Castle Resolve",
    description:
      "Turn one Google-wrapped result link (https://www.google.com/goto?url=…, as printed by castle_cdp_search " +
      "with engine google) into the real destination URL, by following it in the same Chrome. Opens one tab, " +
      "reads where it lands, closes it. It does not read the page. Resolve only the results you mean to open: " +
      "Google answers a burst of links with a \"Verifying your request\" countdown, which this tool reports " +
      "instead of waiting out. Other engines already return real URLs and need no resolving.",
    promptSnippet: "Resolve one Google-wrapped result link to its real URL (castle_cdp_resolve)",
    promptGuidelines: [
      "castle_cdp_search on Google returns wrapped links. Read the titles, snippets and the breadcrumb Google " +
        "shows for each, decide which result you actually need, and resolve only that one with castle_cdp_resolve. " +
        "Do not resolve every result: that burst is what triggers Google's verification page.",
    ],
    parameters: Type.Object({
      url: Type.String({
        description: "A https://www.google.com/goto?url=… link exactly as castle_cdp_search printed it.",
        minLength: 1,
      }),
    }),
    async execute(_toolCallId, params, signal, _onUpdate, _ctx) {
      const timeout = AbortSignal.timeout(SEARCH_TIMEOUT_MS);
      const combined = signal ? AbortSignal.any([signal, timeout]) : timeout;
      const out = await resolveLink(params.url.trim(), combined);
      return {
        content: [{ type: "text", text: `${out.to}\n(resolved via Chrome at ${out.endpoint})` }],
        details: { from: out.from, to: out.to, endpoint: out.endpoint },
      };
    },
  });

  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. Google results come back as wrapped links with the " +
      "breadcrumb Google shows; resolve the ones you need with castle_cdp_resolve. " +
      "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;
      }
    },
  });
}