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;
}
},
});
}
|