@zeldrisho/pi-web-fetch 0.15.0 → 0.16.1

package/src/challenges/aws-waf-solver.ts

@@ -0,0 +1,65 @@
+import { createHash, scryptSync } from "node:crypto";
+
+/** Bounded proof-of-work request sent to the AWS WAF worker. */
+export interface AwsWafWork {
+  type: string;
+  input: string;
+  checksum: string;
+  difficulty: number;
+  memory: number;
+}
+
+/** Returns whether the first `difficulty` bits of a hash are zero. */
+export function hasLeadingZeroBits(hash: Uint8Array, difficulty: number): boolean {
+  for (let bit = 0; bit < difficulty; bit += 1) {
+    const byte = hash[Math.floor(bit / 8)];
+
+    if (byte === undefined || (byte & (1 << (7 - (bit % 8)))) !== 0) return false;
+  }
+
+  return true;
+}
+
+/** Solves the bandwidth challenge using a bounded zero-filled buffer. */
+export function solveAwsWafBandwidth(difficulty: number): string {
+  const sizes = [1_024, 1_024, 10_240, 102_400, 1_048_576, 10_485_760];
+
+  return Buffer.alloc(sizes[Math.min(5, Math.max(0, difficulty))] ?? 1_024).toString("base64");
+}
+
+/** Computes bounded proof-of-work inside a disposable worker. */
+export function solveAwsWafWork(work: AwsWafWork): string | undefined {
+  if (!Number.isInteger(work.difficulty) || work.difficulty < 1 || work.difficulty > 32)
+    return undefined;
+
+  if (work.type === "NetworkBandwidth") return solveAwsWafBandwidth(work.difficulty);
+
+  if (work.type === "SHA256") {
+    for (let nonce = 0; nonce < 2_000_000; nonce += 1) {
+      const digest = createHash("sha256").update(`${work.input}${work.checksum}${nonce}`).digest();
+
+      if (hasLeadingZeroBits(digest, work.difficulty)) return String(nonce);
+    }
+  }
+
+  if (work.type === "HashcashScrypt") {
+    const cost = work.memory > 0 ? work.memory : 128;
+
+    // Reject large or invalid costs before invoking OpenSSL.
+    if (!Number.isInteger(cost) || cost < 2 || cost > 16_384 || (cost & (cost - 1)) !== 0)
+      return undefined;
+
+    for (let nonce = 0; nonce < 100_000; nonce += 1) {
+      const digest = scryptSync(`${work.input}${work.checksum}${nonce}`, work.checksum, 32, {
+        N: cost,
+        r: 8,
+        p: 1,
+        maxmem: 64 * 1_024 * 1_024,
+      });
+
+      if (hasLeadingZeroBits(digest, work.difficulty)) return String(nonce);
+    }
+  }
+
+  return undefined;
+}

package/src/challenges/aws-waf-worker.ts

@@ -0,0 +1,14 @@
+/* v8 ignore file -- exercised through worker integration tests. */
+import { parentPort, workerData } from "node:worker_threads";
+import { solveAwsWafWork, type AwsWafWork } from "./aws-waf-solver.ts";
+
+const work: AwsWafWork = workerData;
+
+try {
+  parentPort?.postMessage({ kind: "solved", solution: solveAwsWafWork(work) });
+} catch (error) {
+  parentPort?.postMessage({
+    kind: "error",
+    message: error instanceof Error ? error.message : "AWS WAF worker failed.",
+  });
+}

package/src/challenges/aws-waf.ts

@@ -0,0 +1,34 @@
+import type { IncomingMessage } from "node:http";
+import type { RedirectDependencies } from "../network/network-redirects.ts";
+
+/** Cookie that carries an AWS WAF challenge token. */
+export const AWS_WAF_TOKEN_COOKIE = "aws-waf-token";
+
+/** A resolved AWS WAF token and the browser identity it was issued for. */
+export interface AwsWafSolution {
+  token: string;
+  userAgent?: string;
+}
+
+/** A request to resolve one AWS WAF challenge. */
+export interface AwsWafResolveContext {
+  /** The page URL withheld by the challenge. */
+  url: URL;
+  /** Cancels the solver request. */
+  signal: AbortSignal;
+  /** Network dependencies used by the resolver. */
+  network?: RedirectDependencies;
+}
+
+/** Resolves an AWS WAF challenge to a token accepted by the target origin. */
+export interface AwsWafResolver {
+  resolve(context: AwsWafResolveContext): Promise<AwsWafSolution | undefined>;
+}
+
+/** Returns whether a response carries an AWS WAF challenge. */
+export function isAwsWafChallengeResponse(response: Pick<IncomingMessage, "headers">): boolean {
+  const value = response.headers["x-amzn-waf-action"];
+  const action = Array.isArray(value) ? value[0] : value;
+
+  return action?.trim().toLowerCase() === "challenge";
+}

package/src/index.ts

@@ -1,5 +1,5 @@
-import { getMarkdownTheme, type ExtensionAPI } from "@earendil-works/pi-coding-agent";
-import { Markdown, Text } from "@earendil-works/pi-tui";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { Text } from "@earendil-works/pi-tui";
 import { Type } from "typebox";
 import {
   FETCH_MAX_OFFSET_CHARACTERS,
@@ -7,7 +7,7 @@ import {
   FETCH_MAX_URL_CHARACTERS,
 } from "./limits";
 import { redactUrlForDisplay } from "./redact";
-import { formatCollapsibleOutput } from "./render";
+import { renderCollapsibleOutput } from "./render";
 import { executeWebFetch, webFetchOutputSchema, type WebFetchDetails } from "./service";
 
 export { ExpiringLruCache } from "./cache";
@@ -20,6 +20,13 @@ export {
   type AnubisSolveContext,
 } from "./challenges/anubis.ts";
 
+export {
+  AWS_WAF_TOKEN_COOKIE,
+  isAwsWafChallengeResponse,
+  type AwsWafResolveContext,
+  type AwsWafResolver,
+} from "./challenges/aws-waf.ts";
+
 export { CookieJar } from "./network/cookies.ts";
 
 export type { FetchResult } from "./document/content.ts";
@@ -101,7 +108,7 @@ export default function (pi: ExtensionAPI) {
     /** Renders the redacted URL, optional focus query, and character offset in the call preview. */
     renderCall(args, theme) {
       return new Text(
-        `${theme.fg("toolTitle", theme.bold("web_fetch"))} ${theme.fg("accent", redactUrlForDisplay(args.url))}`,
+        `${theme.fg("toolTitle", theme.bold("web_fetch"))} ${theme.fg("accent", redactUrlForDisplay(args.url))}${args.query === undefined ? "" : theme.fg("muted", ` query=${JSON.stringify(args.query)}`)}`,
         0,
         0,
       );
@@ -115,15 +122,21 @@ export default function (pi: ExtensionAPI) {
      * Renders partial text or "Fetching…" while pending. Completed text gets a
      * detail summary when available and a collapsible body; absent text shows "No content".
      */
-    renderResult(result, { expanded, isPartial }, theme) {
+    renderResult(result, { expanded, isPartial }, theme, context) {
       const content = result.content.find((item) => item.type === "text");
 
       if (isPartial) return new Text(theme.fg("warning", content?.text || "Fetching…"), 0, 0);
 
       return content?.type === "text"
-        ? new Markdown(formatCollapsibleOutput(content.text, expanded), 0, 0, getMarkdownTheme(), {
-            color: (text) => theme.fg("toolOutput", text),
-          })
+        ? renderCollapsibleOutput(
+            content.text,
+            expanded,
+            theme,
+            context?.isError,
+            result.details?.truncated && result.details.nextOffset !== undefined
+              ? `[Content truncated. Continue with offset=${result.details.nextOffset} to read the next chunk.]`
+              : undefined,
+          )
         : new Text(theme.fg("dim", "No content"), 0, 0);
     },
   });

package/src/network/cookies.ts

@@ -211,6 +211,47 @@ export class CookieJar {
     this.#evict();
   }
 
+  /**
+   * Stores a challenge token returned by a resolver.
+   *
+   * The token is untrusted input. It is accepted only for the target origin,
+   * scoped to the whole origin, and held in memory for the current acquisition.
+   *
+   * @param name - Cookie name
+   * @param value - Cookie value
+   * @param url - The page URL the cookie belongs to
+   * @returns Whether the cookie was accepted
+   */
+  set(name: string, value: string, url: URL): boolean {
+    const token = value.trim();
+
+    if (
+      name.length === 0 ||
+      name.length > MAX_COOKIE_NAME_CHARACTERS ||
+      !COOKIE_NAME.test(name) ||
+      token.length === 0 ||
+      token.length > MAX_COOKIE_VALUE_CHARACTERS ||
+      hasInvalidCookieCharacter(token) ||
+      /[\s";,\\]/u.test(token)
+    ) {
+      return false;
+    }
+
+    const key = `${url.origin}\0${name}\0/`;
+    this.#cookies.delete(key);
+    this.#cookies.set(key, {
+      name,
+      value: token,
+      origin: url.origin,
+      path: "/",
+      secure: url.protocol === "https:",
+      createdAt: Date.now(),
+    });
+    this.#evict();
+
+    return true;
+  }
+
   /**
    * Returns a Cookie header value for a URL, excluding expired, insecure, and
    * cross-origin cookies.

package/src/network/network-redirects.ts

@@ -29,11 +29,16 @@ export interface RedirectDependencies {
     target: ValidatedTarget,
     signal: AbortSignal,
     headers?: Readonly<Record<string, string>>,
+    options?: { method?: string; body?: string | Uint8Array },
   ) => Promise<IncomingMessage>;
   random?: () => number;
   sleep?: (milliseconds: number, signal: AbortSignal) => Promise<void>;
   /** Origin-scoped cookies for the current acquisition, including challenge handshakes. */
   cookies?: CookieJar;
+  /** User-Agent override that persists across redirects. */
+  userAgent?: string;
+  /** Headers that survive a redirect, used by challenge handshakes that need browser identity. */
+  persistentHeaders?: Readonly<Record<string, string>>;
 }
 
 /** Creates the consistent abort error used by redirect and retry coordination. */
@@ -144,7 +149,14 @@ async function coordinatedRequest(
     const release = await acquireOrigin(target.url.origin, signal);
     let response: IncomingMessage;
     const cookieHeader = dependencies.cookies?.headerFor(target.url);
-    const requestHeaders = cookieHeader ? { Cookie: cookieHeader, ...headers } : headers;
+
+    const headersWithUserAgent = dependencies.userAgent
+      ? { "User-Agent": dependencies.userAgent, ...headers }
+      : headers;
+
+    const requestHeaders = cookieHeader
+      ? { Cookie: cookieHeader, ...headersWithUserAgent }
+      : headersWithUserAgent;
 
     try {
       response = await request(target, signal, requestHeaders);
@@ -197,7 +209,8 @@ export async function requestFollowingRedirects(
     response.resume();
     target = await awaitWithAbort(validateUrl(new URL(location, target.url)), signal);
     // Validators describe the originally cached resource and must not leak to a redirect target.
-    requestHeaders = {};
+    // Challenge handshakes opt into a fixed browser identity instead.
+    requestHeaders = dependencies.persistentHeaders ?? {};
   }
 
   throw new Error("web_fetch followed too many redirects.");

package/src/network/network-transport.ts

@@ -8,6 +8,9 @@ import type { ValidatedTarget } from "./network-policy.ts";
 /** Maximum size in bytes for raw HTTP response bodies. */
 export const FETCH_MAX_BYTES = 5 * 1_024 * 1_024;
 
+/** Fixed User-Agent used for page requests and challenge solvers. */
+export const WEB_FETCH_USER_AGENT = "Mozilla/5.0 (compatible; PiWebFetch/1.0; +https://pi.dev)";
+
 /**
  * Per-address connect deadline. A hanging address is abandoned after this
  * budget so the next validated address can be tried inside the overall
@@ -67,6 +70,8 @@ async function requestOnce(
   signal: AbortSignal,
   attemptTimeoutMs: number,
   headers: Readonly<Record<string, string>>,
+  method = "GET",
+  body?: string | Uint8Array,
 ): Promise<IncomingMessage> {
   const family = isIP(address);
 
@@ -114,11 +119,12 @@ async function requestOnce(
       target.url,
       {
         lookup,
+        method,
         signal: controller.signal,
         headers: {
           Accept:
             "text/markdown;q=1.0, text/x-markdown;q=0.9, text/plain;q=0.8, text/html;q=0.7, application/xhtml+xml;q=0.7, application/json;q=0.7, application/xml;q=0.7, */*;q=0.1",
-          "User-Agent": "Mozilla/5.0 (compatible; PiWebFetch/1.0; +https://pi.dev)",
+          "User-Agent": WEB_FETCH_USER_AGENT,
           ...headers,
         },
       },
@@ -139,7 +145,7 @@ async function requestOnce(
       if (attemptExpired) finish(() => reject(unreachableError()));
       else finish(() => reject(error));
     });
-    outgoing.end();
+    outgoing.end(body);
   });
 }
 
@@ -148,6 +154,10 @@ export interface RequestPinnedOptions {
   attemptTimeoutMs?: number;
   /** Additional fixed request headers, used for conditional cache revalidation. */
   headers?: Readonly<Record<string, string>>;
+  /** HTTP method; defaults to GET. */
+  method?: string;
+  /** Request body for non-GET requests. */
+  body?: string | Uint8Array;
 }
 
 /**
@@ -169,7 +179,15 @@ export async function requestPinned(
 
   for (const address of addresses) {
     try {
-      return await requestOnce(target, address, signal, attemptTimeoutMs, options.headers ?? {});
+      return await requestOnce(
+        target,
+        address,
+        signal,
+        attemptTimeoutMs,
+        options.headers ?? {},
+        options.method,
+        options.body,
+      );
     } catch (error) {
       lastError = error;
 

package/src/render.ts

@@ -1,21 +1,46 @@
-import { keyHint } from "@earendil-works/pi-coding-agent";
+import { getMarkdownTheme, keyHint } from "@earendil-works/pi-coding-agent";
+import { Container, Markdown, Spacer, Text } from "@earendil-works/pi-tui";
 
 const COLLAPSED_LINES = 10;
 
-/** Format tool output using Pi's built-in collapsed-preview convention. */
-export function formatCollapsibleOutput(output: string, expanded: boolean): string {
+/** Render document content separately from Pi-style interface notices. */
+export function renderCollapsibleOutput(
+  output: string,
+  expanded: boolean,
+  theme: { fg(color: "toolOutput" | "muted" | "warning", text: string): string },
+  isError = false,
+  warning?: string,
+): Container {
   const lines = output.replace(/\r\n/g, "\n").split("\n");
 
   while (lines.at(-1) === "") lines.pop();
-  const totalLines = lines.length;
-  const visibleLines = expanded ? lines : lines.slice(0, COLLAPSED_LINES);
-  let text = visibleLines.join("\n");
+  const visible = expanded ? lines : lines.slice(0, COLLAPSED_LINES);
+  const remaining = lines.length - visible.length;
+  const component = new Container();
+  const body = `\n${visible.join("\n")}`;
 
-  const remaining = totalLines - visibleLines.length;
+  component.addChild(
+    isError
+      ? new Text(theme.fg("toolOutput", body), 0, 0)
+      : new Markdown(body, 0, 0, getMarkdownTheme(), {
+          color: (text) => theme.fg("toolOutput", text),
+        }),
+  );
 
   if (remaining > 0) {
-    text += `\n... (${remaining} more lines, ${totalLines} total, ${keyHint("app.tools.expand", "to expand")})`;
+    component.addChild(
+      new Text(
+        `${theme.fg("muted", `... (${remaining} more lines, ${lines.length} total,`)} ${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`,
+        0,
+        0,
+      ),
+    );
   }
 
-  return `\n${text}`;
+  if (warning) {
+    component.addChild(new Spacer(1));
+    component.addChild(new Text(theme.fg("warning", warning), 0, 0));
+  }
+
+  return component;
 }

package/src/service.ts

@@ -8,7 +8,7 @@ import {
   truncateHead,
   type AgentToolResult,
 } from "@earendil-works/pi-coding-agent";
-import { ExpiringLruCache, stableKeyHash, type CachePersistence } from "./cache";
+import { ExpiringLruCache, stableKeyHash, canonicalKey, type CachePersistence } from "./cache";
 import { sliceCompleteDocument, type CompleteDocument } from "./document/content.ts";
 import { revalidateCompleteDocument, type FetchRemoteDependencies } from "./acquisition/fetch.ts";
 import { fetchDocumentWithLlmsTxtSupport } from "./acquisition/llms.ts";
@@ -430,7 +430,8 @@ export async function executeWebFetch(
 
   const displayRequestedUrl = redactUrlForDisplay(params.url);
   const sourceUrl = urlWithoutFragment(params.url);
-  let document = fetchCache.get(sourceUrl);
+  const cacheKey = canonicalKey({ url: sourceUrl });
+  let document = fetchCache.get(cacheKey);
   const now = Date.now();
 
   const fresh = Boolean(
@@ -459,7 +460,7 @@ export async function executeWebFetch(
     const stale = document;
 
     const acquisition = await inflightFetches.run(
-      sourceUrl,
+      cacheKey,
       async (sharedSignal) => {
         let fetched: CompleteDocument;
         let acquisitionStatus: FetchAcquisition["cacheStatus"] = "miss";
@@ -479,7 +480,7 @@ export async function executeWebFetch(
         }
 
         if (!fetched.extractionDiagnostics?.botWall) {
-          fetchCache.set(sourceUrl, fetched, Date.now() + CACHE_STALE_RETENTION_MS);
+          fetchCache.set(cacheKey, fetched, Date.now() + CACHE_STALE_RETENTION_MS);
         }
 
         return { document: fetched, cacheStatus: acquisitionStatus };

package/workers/aws-waf-solver.js

@@ -0,0 +1,56 @@
+import { createHash, scryptSync } from "node:crypto";
+
+/** Returns whether the first `difficulty` bits of a hash are zero. */
+export function hasLeadingZeroBits(hash, difficulty) {
+  for (let bit = 0; bit < difficulty; bit += 1) {
+    const byte = hash[Math.floor(bit / 8)];
+
+    if (byte === undefined || (byte & (1 << (7 - (bit % 8)))) !== 0) return false;
+  }
+
+  return true;
+}
+
+/** Solves the bandwidth challenge using a bounded zero-filled buffer. */
+export function solveAwsWafBandwidth(difficulty) {
+  const sizes = [1_024, 1_024, 10_240, 102_400, 1_048_576, 10_485_760];
+
+  return Buffer.alloc(sizes[Math.min(5, Math.max(0, difficulty))] ?? 1_024).toString("base64");
+}
+
+/** Computes bounded proof-of-work inside a disposable worker. */
+export function solveAwsWafWork(work) {
+  if (!Number.isInteger(work.difficulty) || work.difficulty < 1 || work.difficulty > 32)
+    return undefined;
+
+  if (work.type === "NetworkBandwidth") return solveAwsWafBandwidth(work.difficulty);
+
+  if (work.type === "SHA256") {
+    for (let nonce = 0; nonce < 2_000_000; nonce += 1) {
+      const digest = createHash("sha256").update(`${work.input}${work.checksum}${nonce}`).digest();
+
+      if (hasLeadingZeroBits(digest, work.difficulty)) return String(nonce);
+    }
+  }
+
+  if (work.type === "HashcashScrypt") {
+    const cost = work.memory > 0 ? work.memory : 128;
+
+    // Reject large or invalid costs before invoking OpenSSL.
+    if (!Number.isInteger(cost) || cost < 2 || cost > 16_384 || (cost & (cost - 1)) !== 0)
+      return undefined;
+
+    for (let nonce = 0; nonce < 100_000; nonce += 1) {
+      const digest = scryptSync(`${work.input}${work.checksum}${nonce}`, work.checksum, 32, {
+        N: cost,
+        r: 8,
+        p: 1,
+        maxmem: 64 * 1_024 * 1_024,
+      });
+
+      if (hasLeadingZeroBits(digest, work.difficulty)) return String(nonce);
+    }
+  }
+
+  return undefined;
+}

package/workers/aws-waf-worker.js

@@ -0,0 +1,14 @@
+/* v8 ignore file -- exercised through worker integration tests. */
+import { parentPort, workerData } from "node:worker_threads";
+import { solveAwsWafWork } from "./aws-waf-solver.js";
+
+const work = workerData;
+
+try {
+  parentPort?.postMessage({ kind: "solved", solution: solveAwsWafWork(work) });
+} catch (error) {
+  parentPort?.postMessage({
+    kind: "error",
+    message: error instanceof Error ? error.message : "AWS WAF worker failed.",
+  });
+}