@lincoln504/pi-research 1.4.1 → 1.5.0

package/src/utils/tool-usage-tracker.ts

@@ -195,12 +195,15 @@ export class ToolUsageTracker {
 }
 
 /**
- * Create default tool limits for a researcher
+ * Create default tool limits for a researcher.
+ *
+ * @param cachingActive - Whether prompt caching is active for this run's
+ * resolved model; see getMaxScrapeBatches in constants.ts.
  */
-export function createDefaultToolLimits(config?: Config): ToolLimits {
+export function createDefaultToolLimits(config?: Config, cachingActive = false): ToolLimits {
   return {
     gathering: getMaxGatheringCalls(config),
-    scrape: getMaxScrapeBatches(config),
+    scrape: getMaxScrapeBatches(config, cachingActive),
     search: 1,
     youtube_transcript: 1,
     grep: MAX_GREP_CALLS,

package/src/web-research/browser-search.ts

@@ -5,7 +5,7 @@
  */
 
 import { runWorkerSearch } from '../infrastructure/browser/task-execution-service.ts';
-import { getMaxWorkers } from '../infrastructure/browser/config.ts';
+import { getMaxWorkers, COLD_START_ALLOWANCE_MS } from '../infrastructure/browser/config.ts';
 import { isTaskTimeoutError } from '../infrastructure/browser/browser-error-utils.ts';
 import { logger } from '../logger.ts';
 import { safeUnref } from '../utils/safe-unref.ts';
@@ -82,8 +82,9 @@ export async function performSearch(
     // cannot stall the entire Promise.all burst.
     //
     // The budget mirrors the scheduler's own task ceiling in
-    // BrowserTaskScheduler.runSearch (= SEARCH_TIMEOUT_MS + BROWSER_TASK_TIMEOUT_MS)
-    // so this layer never preempts the worker's deadline with a shorter one.
+    // BrowserTaskScheduler.runSearch (= SEARCH_TIMEOUT_MS + BROWSER_TASK_TIMEOUT_MS +
+    // COLD_START_ALLOWANCE_MS) so this layer never preempts the worker's deadline with
+    // a shorter one.
     //
     // Both clocks now start at DISPATCH (see startQueryDeadline below), which is
     // what makes mirroring the value correct. While this one started at enqueue,
@@ -95,7 +96,13 @@ export async function performSearch(
     // query's worth of wait, and a burst four times the pool's reach exhausts any
     // constant. Measuring execution instead of waiting is the fix; the sum
     // survives as the budget because a search legitimately costs nav plus overhead.
-    const QUERY_TIMEOUT_MS = (config?.SEARCH_TIMEOUT_MS ?? 45_000) + (config?.BROWSER_TASK_TIMEOUT_MS ?? 10_000);
+    //
+    // COLD_START_ALLOWANCE_MS is part of the sum for the same reason it is part of
+    // the scheduler's own ceiling: the first query dispatched to a freshly created or
+    // just-reset worker pays a real browser launch + context creation inline before
+    // any navigation starts, and this layer's budget must not fire before the
+    // scheduler's own (correct, larger) deadline gets a chance to.
+    const QUERY_TIMEOUT_MS = (config?.SEARCH_TIMEOUT_MS ?? 45_000) + (config?.BROWSER_TASK_TIMEOUT_MS ?? 10_000) + COLD_START_ALLOWANCE_MS;
 
     // Deduplicated: resultMap and the failures map are both keyed by the query
     // string, so two tasks for the same string race last-writer-wins — instance A

package/src/web-research/web-scraper.ts

@@ -75,6 +75,15 @@ async function convertToMarkdown(html: string): Promise<string> {
   return converter(html);
 }
 
+// TODO(hardening, medium): this synchronous WASM parse runs on the MAIN
+// process's event loop — a large/complex PDF can block it for seconds,
+// starving every concurrent scrape, the TUI render loop, and timers. The
+// correct fix is a worker_threads-based extraction (mirroring the browser
+// pool's CPU isolation pattern), but it is deliberately NOT done in this
+// pass: pdf-oxide-wasm is dynamically imported here and unit tests mock the
+// module via vitest, which cannot reach into a real worker thread — moving
+// the parse requires a companion test-seam redesign first. Bounded today by
+// MAX_PDF_SIZE (100MB) and the scrape task's own deadline.
 async function extractPdfToMarkdown(bytes: Uint8Array): Promise<string> {
   if (bytes.length > MAX_PDF_SIZE) {
     const sizeMB = Math.round(bytes.length / 1024 / 1024);
@@ -123,23 +132,34 @@ async function extractPdfToMarkdown(bytes: Uint8Array): Promise<string> {
 
   try {
     const doc = new WasmPdfDocument(bytes);
-    const pageCount = doc.pageCount();
-    
-    let markdown = `# PDF Document\n\n**Pages:** ${pageCount}\n\n`;
-    
+    // free() is native/WASM-side memory, not GC-reclaimed — it must run on
+    // every exit, including the case toMarkdownAll() AND the per-page fallback
+    // both throw (a malformed/encrypted PDF the codebase already treats as
+    // routine on the open web, so this path fires regularly). It used to sit
+    // after the fallback loop, reachable only when nothing below it threw —
+    // the one case that needed the fallback landing squarely on the one case
+    // that skipped the free(), leaking one WasmPdfDocument per occurrence in a
+    // browser-server process that stays warm across many research sessions.
     try {
-      markdown += doc.toMarkdownAll();
-    } catch {
-      for (let i = 0; i < pageCount; i++) {
-        markdown += `## Page ${i + 1}\n\n${doc.toMarkdown(i)}\n\n`;
+      const pageCount = doc.pageCount();
+
+      let markdown = `# PDF Document\n\n**Pages:** ${pageCount}\n\n`;
+
+      try {
+        markdown += doc.toMarkdownAll();
+      } catch {
+        for (let i = 0; i < pageCount; i++) {
+          markdown += `## Page ${i + 1}\n\n${doc.toMarkdown(i)}\n\n`;
+        }
       }
+
+      const pdfExtractionDuration = Date.now() - pdfExtractionStart;
+      metrics.observe('scrape_pdf_conversion_ms', pdfExtractionDuration);
+      metrics.increment('scrape_pdf_conversions_total', 1, { status: 'success', pages: String(pageCount) });
+      return markdown;
+    } finally {
+      doc.free();
     }
-    
-    doc.free();
-    const pdfExtractionDuration = Date.now() - pdfExtractionStart;
-    metrics.observe('scrape_pdf_conversion_ms', pdfExtractionDuration);
-    metrics.increment('scrape_pdf_conversions_total', 1, { status: 'success', pages: String(pageCount) });
-    return markdown;
   } catch (e) {
     const msg = e instanceof Error ? e.message : String(e);
     // debug, not error: a malformed/encrypted/scanned-image-only PDF is an
@@ -300,7 +320,7 @@ async function scrapeWithFetch(url: string, signal?: AbortSignal): Promise<Scrap
     const contentLength = response.headers.get('content-length');
     if (contentLength) {
       const size = parseInt(contentLength, 10);
-      if (contentType.includes('application/pdf') || url.toLowerCase().endsWith('.pdf')) {
+      if (contentType.includes('application/pdf') || currentUrl.toLowerCase().endsWith('.pdf')) {
         if (size > MAX_PDF_SIZE) {
           const sizeMB = Math.round(size / 1024 / 1024);
           logger.warn(`[Scrapers] PDF too large (Content-Length: ${sizeMB}MB, max 100MB), skipping`);
@@ -319,7 +339,7 @@ async function scrapeWithFetch(url: string, signal?: AbortSignal): Promise<Scrap
       }
     }
     
-    if (contentType.includes('application/pdf') || url.toLowerCase().endsWith('.pdf')) {
+    if (contentType.includes('application/pdf') || currentUrl.toLowerCase().endsWith('.pdf')) {
       let pdfBytes: Uint8Array;
       try {
         pdfBytes = await readBodyCapped(response, MAX_PDF_SIZE);

package/src/youtube/index.ts

@@ -10,6 +10,8 @@ import type { AgentToolResult } from '@earendil-works/pi-coding-agent';
 import { selectVideoIds } from './video-id.ts';
 import { fetchVideoTranscripts, type VideoTranscript } from './transcript-client.ts';
 import { metrics } from '../utils/metrics.ts';
+import { truncateWithMarker } from '../utils/text-utils.ts';
+import { MAX_SCRAPE_CONTENT_CHARS_PER_DOC } from '../constants.ts';
 
 export interface YoutubeTranscriptCommandOptions {
   urls: string[];
@@ -131,7 +133,15 @@ function formatResults(results: VideoTranscript[], rejected: string[]): string {
     // model can copy verbatim; omitting the line entirely for untitled videos
     // left them without any prompt-visible citation string to reuse.
     md += `**Cite as:** ${r.url}${r.title ? ` — '${r.title}'${r.author ? ` by ${r.author}` : ''}` : ''}\n`;
-    md += `\n${r.text}\n\n---\n\n`;
+    // Per-document cap on content entering LLM context — parity with the
+    // scrape tool's own MAX_SCRAPE_CONTENT_CHARS_PER_DOC guard. The raw HTTP
+    // body is capped at 32MB (readTextCapped) and the knowledge-store
+    // ingestion path truncates too, but this tool's direct response to the
+    // calling agent had no cap of its own: a single long/dense-captioned
+    // video (a multi-hour livestream VOD, worst case with maxVideos up to 5)
+    // could inject an unbounded, multi-megabyte payload straight into the
+    // researcher's context.
+    md += `\n${truncateWithMarker(r.text ?? '', MAX_SCRAPE_CONTENT_CHARS_PER_DOC)}\n\n---\n\n`;
   }
 
   if (failed.length > 0) {

package/src/youtube/potoken.ts

@@ -44,6 +44,7 @@
 import { logger } from '../logger.ts';
 import { safeUnref } from '../utils/safe-unref.ts';
 import { metrics } from '../utils/metrics.ts';
+import { readJsonCapped } from '../utils/http-body.ts';
 
 /**
  * Well-known request key for the YouTube web PoToken program. This is a public
@@ -229,7 +230,13 @@ async function doMint(
       body: JSON.stringify([requestKey, botguardResponse]),
       signal: opts.signal,
     });
-    const integrityJson = (await integrityResponse.json()) as unknown[];
+    if (!integrityResponse.ok) {
+      throw new Error(`HTTP ${integrityResponse.status} from attestation server`);
+    }
+    // The attestation response is a short JSON array; readJsonCapped bounds it the
+    // same way as every other externally-controlled body this codebase reads (the
+    // default 32MB cap is already generous for this endpoint).
+    const integrityJson = await readJsonCapped<unknown[]>(integrityResponse);
     const integrityToken = integrityJson?.[0] as string | undefined;
     if (!integrityToken) {
       throw new Error('Integrity token rejected by attestation server (likely a non-residential IP)');

package/src/youtube/transcript-client.ts

@@ -18,6 +18,7 @@
  */
 
 import { withTimeout, retryWithBackoff, createTimeoutSignal } from '../web-research/retry-utils.ts';
+import { getSsrfSafeFetcher } from '../web-research/scraper-utils.ts';
 import { CircuitBreaker } from '../utils/circuit-breaker.ts';
 import { logger } from '../logger.ts';
 import { redactSecrets } from '../utils/log-utils.ts';
@@ -109,7 +110,31 @@ export async function fetchVideoTranscripts(
   } catch {
     // non-fatal: logging stays at default if the API ever changes
   }
-  const fetchImpl: typeof fetch = (input, init) => fetch(input, init);
+  // Connect-time SSRF pin (see scraper-utils.ts's getSsrfSafeFetcher): closes the
+  // gap between assertTrustedTimedtextUrl's host check below and fetch's own
+  // re-resolution at connect time. The host allowlist only constrains WHICH
+  // hostname a redirect can name; it says nothing about what that hostname
+  // resolves to at the moment the socket opens, and this codebase already
+  // treats a resolver answering differently between validation and connect as
+  // a real threat elsewhere (thread-worker-messaging.ts's DNS-rebinding
+  // backstop applies even to already-trusted hosts). One fetchImpl closure
+  // covers every network call in this batch: the BotGuard/attestation fetches
+  // inside mintPoTokens() and the timedtext fetch below. Falls back to the
+  // plain global fetch for non-string/URL input (never seen in practice here)
+  // since the dispatcher pairing requires a URL string; see getSsrfSafeFetcher's
+  // own doc comment for why the fetch and its dispatcher must travel together.
+  const ssrfFetcher = await getSsrfSafeFetcher();
+  const fetchImpl: typeof fetch = ssrfFetcher
+    ? (input, init) => {
+        if (typeof input === 'string' || input instanceof URL) {
+          return ssrfFetcher.fetch(String(input), {
+            ...(init ?? {}),
+            dispatcher: ssrfFetcher.dispatcher,
+          }) as unknown as Promise<Response>;
+        }
+        return fetch(input, init);
+      }
+    : (input, init) => fetch(input, init);
 
   // 1) Throwaway session to obtain stable visitor data for the PoToken binding.
   let visitorData: string;
@@ -245,10 +270,36 @@ async function fetchOne(
           // response headers AND the body read are bounded and tear the socket
           // down on timeout/cancellation (withTimeout alone would not bound text()).
           const fetchSignal = createTimeoutSignal(timeoutMs, signal);
-          const response = await fetchImpl(url, {
-            signal: fetchSignal,
-            headers: { 'User-Agent': 'Mozilla/5.0', 'Accept-Language': lang },
-          });
+          // Manual redirect following, re-validating each hop against the same
+          // trusted-host allowlist as the first hop. `redirect: 'follow' (the
+          // default) would carry the PoToken-bearing request to whatever host a
+          // 3xx Location header names with no revalidation — assertTrustedTimedtextUrl
+          // above only ever checked base_url, the pre-redirect URL. fetchImpl
+          // carries the connect-time IP pin (see its construction above), so each
+          // hop is also protected against a resolver answering differently
+          // between this host check and the moment the socket actually opens.
+          const MAX_REDIRECTS = 5;
+          let currentUrl = url;
+          let response!: Response;
+          for (let hop = 0; hop <= MAX_REDIRECTS; hop++) {
+            response = await fetchImpl(currentUrl, {
+              signal: fetchSignal,
+              redirect: 'manual',
+              headers: { 'User-Agent': 'Mozilla/5.0', 'Accept-Language': lang },
+            });
+            if (response.status >= 300 && response.status < 400) {
+              const location = response.headers.get('location');
+              if (!location) break; // No Location header — treat as final response
+              const resolved = new URL(location, currentUrl).href;
+              assertTrustedTimedtextUrl(resolved);
+              // Drain the redirect response before following: an unconsumed body
+              // pins its socket (undici keeps it open until read/cancelled or GC).
+              void response.body?.cancel()?.catch(() => { /* best-effort */ });
+              currentUrl = resolved;
+              continue;
+            }
+            break; // Non-redirect response
+          }
           // Surface the real HTTP status. Without this, a 403 (age-restricted) or
           // 404 (removed) returns HTML that parseJson3 yields '' for, which would
           // be misreported below as a bot-protection/PoToken failure.