@tonyclaw/agent-inspector 3.1.20 → 3.1.21

package/src/lib/resourceLimits.ts

@@ -11,6 +11,197 @@ export const MAX_ACTIVE_PROVIDER_TESTS = 2;
 export const MAX_PROVIDER_TEST_CONCURRENCY = 4;
 export const MAX_ACTIVE_PROXY_REQUESTS = 16;
 
+const positiveInteger = z.number().int().positive();
+
+export const RuntimeBudgetPolicySchema = z
+  .object({
+    proxy: z.object({
+      requestBytes: positiveInteger,
+      responseBytes: positiveInteger,
+      headerBytes: positiveInteger,
+      deadlineMs: positiveInteger,
+      idleTimeoutMs: positiveInteger,
+      activeRequests: positiveInteger,
+    }),
+    control: z.object({
+      requestBytes: positiveInteger,
+      responseBytes: positiveInteger,
+      headerBytes: positiveInteger,
+      deadlineMs: positiveInteger,
+      idleTimeoutMs: positiveInteger,
+      activeRequests: positiveInteger,
+    }),
+    upstream: z.object({
+      fetchResponseBytes: positiveInteger,
+      fetchDeadlineMs: positiveInteger,
+      replayResponseBytes: positiveInteger,
+      replayDeadlineMs: positiveInteger,
+      metadataResponseBytes: positiveInteger,
+      metadataDeadlineMs: positiveInteger,
+      connectTimeoutMs: positiveInteger,
+      idleTimeoutMs: positiveInteger,
+    }),
+    pagination: z.object({
+      restPageSize: positiveInteger,
+      searchQueryChars: positiveInteger,
+      searchProjectChars: positiveInteger,
+    }),
+    sse: z.object({
+      clients: positiveInteger,
+      initLogs: positiveInteger,
+      heartbeatMs: positiveInteger,
+    }),
+    imports: z.object({
+      archiveBytes: positiveInteger,
+      entries: positiveInteger,
+      entryBytes: positiveInteger,
+      uncompressedBytes: positiveInteger,
+      expansionRatio: positiveInteger,
+    }),
+    fanOut: z.object({
+      providerModels: positiveInteger,
+      providerModelMetadata: positiveInteger,
+      providerTestConcurrency: positiveInteger,
+      activeProviderTests: positiveInteger,
+    }),
+    queues: z.object({
+      proxyRequests: positiveInteger,
+      finalizerJobs: positiveInteger,
+      rawStreamMessages: positiveInteger,
+      rawStreamBytes: positiveInteger,
+      backgroundWorkers: positiveInteger,
+    }),
+  })
+  .superRefine((policy, context) => {
+    if (policy.control.requestBytes > policy.proxy.requestBytes) {
+      context.addIssue({
+        code: "custom",
+        path: ["control", "requestBytes"],
+        message: "control request budget must not exceed proxy request budget",
+      });
+    }
+    if (policy.control.responseBytes > policy.proxy.responseBytes) {
+      context.addIssue({
+        code: "custom",
+        path: ["control", "responseBytes"],
+        message: "control response budget must not exceed proxy response budget",
+      });
+    }
+    if (policy.control.activeRequests > policy.proxy.activeRequests) {
+      context.addIssue({
+        code: "custom",
+        path: ["control", "activeRequests"],
+        message: "control active request budget must not exceed proxy active request budget",
+      });
+    }
+    if (policy.upstream.fetchResponseBytes > policy.proxy.responseBytes) {
+      context.addIssue({
+        code: "custom",
+        path: ["upstream", "fetchResponseBytes"],
+        message: "upstream fetch response budget must not exceed proxy response budget",
+      });
+    }
+    if (policy.upstream.replayResponseBytes > policy.proxy.responseBytes) {
+      context.addIssue({
+        code: "custom",
+        path: ["upstream", "replayResponseBytes"],
+        message: "upstream replay response budget must not exceed proxy response budget",
+      });
+    }
+    if (policy.upstream.metadataResponseBytes > policy.control.responseBytes) {
+      context.addIssue({
+        code: "custom",
+        path: ["upstream", "metadataResponseBytes"],
+        message: "upstream metadata response budget must not exceed control response budget",
+      });
+    }
+    if (policy.upstream.metadataDeadlineMs > policy.control.deadlineMs) {
+      context.addIssue({
+        code: "custom",
+        path: ["upstream", "metadataDeadlineMs"],
+        message: "upstream metadata deadline must not exceed control deadline",
+      });
+    }
+    if (policy.queues.proxyRequests < policy.proxy.activeRequests) {
+      context.addIssue({
+        code: "custom",
+        path: ["queues", "proxyRequests"],
+        message: "proxy request queue budget must cover active proxy requests",
+      });
+    }
+  });
+
+export type RuntimeBudgetPolicy = z.infer<typeof RuntimeBudgetPolicySchema>;
+
+export const DEFAULT_RUNTIME_BUDGET_POLICY = RuntimeBudgetPolicySchema.parse({
+  proxy: {
+    requestBytes: 64 * 1024 * 1024,
+    responseBytes: 256 * 1024 * 1024,
+    headerBytes: 64 * 1024,
+    deadlineMs: 120_000,
+    idleTimeoutMs: 30_000,
+    activeRequests: MAX_ACTIVE_PROXY_REQUESTS,
+  },
+  control: {
+    requestBytes: 8 * 1024 * 1024,
+    responseBytes: 64 * 1024 * 1024,
+    headerBytes: 32 * 1024,
+    deadlineMs: 30_000,
+    idleTimeoutMs: 10_000,
+    activeRequests: 8,
+  },
+  upstream: {
+    fetchResponseBytes: 256 * 1024 * 1024,
+    fetchDeadlineMs: 120_000,
+    replayResponseBytes: 256 * 1024 * 1024,
+    replayDeadlineMs: 120_000,
+    metadataResponseBytes: 10 * 1024 * 1024,
+    metadataDeadlineMs: 30_000,
+    connectTimeoutMs: 10_000,
+    idleTimeoutMs: 30_000,
+  },
+  pagination: {
+    restPageSize: MAX_REST_LOG_PAGE_SIZE,
+    searchQueryChars: MAX_SEARCH_QUERY_CHARS,
+    searchProjectChars: MAX_SEARCH_PROJECT_CHARS,
+  },
+  sse: {
+    clients: MAX_SSE_CLIENTS,
+    initLogs: 500,
+    heartbeatMs: 15_000,
+  },
+  imports: {
+    archiveBytes: 64 * 1024 * 1024,
+    entries: 5_000,
+    entryBytes: 32 * 1024 * 1024,
+    uncompressedBytes: 256 * 1024 * 1024,
+    expansionRatio: 20,
+  },
+  fanOut: {
+    providerModels: MAX_PROVIDER_MODELS,
+    providerModelMetadata: MAX_PROVIDER_MODEL_METADATA,
+    providerTestConcurrency: MAX_PROVIDER_TEST_CONCURRENCY,
+    activeProviderTests: MAX_ACTIVE_PROVIDER_TESTS,
+  },
+  queues: {
+    proxyRequests: MAX_ACTIVE_PROXY_REQUESTS,
+    finalizerJobs: 1_000,
+    rawStreamMessages: 16_384,
+    rawStreamBytes: 64 * 1024 * 1024,
+    backgroundWorkers: 4,
+  },
+});
+
+export function validateRuntimeBudgetPolicy(input: unknown): RuntimeBudgetPolicy {
+  return RuntimeBudgetPolicySchema.parse(input);
+}
+
+export function runtimeBudgetPolicySnapshot(
+  policy: RuntimeBudgetPolicy = DEFAULT_RUNTIME_BUDGET_POLICY,
+): RuntimeBudgetPolicy {
+  return RuntimeBudgetPolicySchema.parse(policy);
+}
+
 export const ResourceLimitErrorCodeSchema = z.enum([
   "invalid-page-request",
   "search-limit-exceeded",

package/src/lib/safeFetch.ts

@@ -1,5 +1,8 @@
 import { lookup } from "node:dns/promises";
+import { request as requestHttp } from "node:http";
+import { request as requestHttps } from "node:https";
 import { isIP } from "node:net";
+import type { IncomingMessage, RequestOptions } from "node:http";
 import { evaluateUpstreamUrl, isInternalHost, type SsrfGuardOptions } from "./ssrfGuard";
 
 export type ResolvedAddress = {
@@ -11,11 +14,18 @@ export type HostResolver = (hostname: string) => Promise<readonly ResolvedAddres
 
 export type SafeFetchOptions = SsrfGuardOptions & {
   resolver?: HostResolver;
+  deadlineMs?: number;
+  responseBytes?: number;
+  pinDns?: boolean;
 };
 
 export type SafeFetchResult =
   | { ok: true; response: Response }
-  | { ok: false; kind: "network-policy" | "redirect"; reason: string };
+  | {
+      ok: false;
+      kind: "network-policy" | "redirect" | "resource-limit" | "timeout";
+      reason: string;
+    };
 
 async function defaultResolver(hostname: string): Promise<readonly ResolvedAddress[]> {
   return lookup(hostname, { all: true, verbatim: true });
@@ -35,12 +45,28 @@ export async function evaluateResolvedUpstreamUrl(
   rawUrl: string,
   options: SafeFetchOptions = {},
 ): Promise<{ allow: true } | { allow: false; reason: string }> {
+  const decision = await evaluateResolvedUpstreamUrlForFetch(rawUrl, options);
+  return decision.allow ? { allow: true } : { allow: false, reason: decision.reason };
+}
+
+type ResolvedUpstreamUrlDecision =
+  | { allow: true; hostname: string; addresses: readonly ResolvedAddress[] }
+  | { allow: false; reason: string };
+
+async function evaluateResolvedUpstreamUrlForFetch(
+  rawUrl: string,
+  options: SafeFetchOptions,
+): Promise<ResolvedUpstreamUrlDecision> {
   const initial = evaluateUpstreamUrl(rawUrl, options);
-  if (!initial.allow || options.bypass === true) return initial;
+  if (!initial.allow) return initial;
 
   const parsed = new URL(rawUrl);
   const hostname = parsed.hostname.replace(/^\[|]$/g, "");
-  if (isIP(hostname) !== 0 || hostname.toLowerCase() === "localhost") return initial;
+  if (options.bypass === true) return { allow: true, hostname, addresses: [] };
+
+  const ipFamily = isIP(hostname);
+  if (ipFamily !== 0)
+    return { allow: true, hostname, addresses: [{ address: hostname, family: ipFamily }] };
 
   let addresses: readonly ResolvedAddress[];
   try {
@@ -65,36 +91,418 @@ export async function evaluateResolvedUpstreamUrl(
     };
   }
 
-  return { allow: true };
+  return { allow: true, hostname, addresses };
 }
 
 function isRedirectStatus(status: number): boolean {
   return status >= 300 && status < 400;
 }
 
+function formatByteLimit(limit: number): string {
+  return `${String(limit)} bytes`;
+}
+
+function parseContentLength(value: string | null): number | null {
+  if (value === null) return null;
+  const parsed = Number(value);
+  if (!Number.isSafeInteger(parsed) || parsed < 0) return null;
+  return parsed;
+}
+
+function isPinnedDnsEnabled(options: SafeFetchOptions): boolean {
+  const disabledForTest =
+    typeof process !== "undefined" &&
+    process.env["NODE_ENV"] === "test" &&
+    process.env["AGENT_INSPECTOR_DISABLE_DNS_PINNING_FOR_TESTS"] === "1";
+  if (disabledForTest) return false;
+  return options.pinDns !== false && options.bypass !== true;
+}
+
+function choosePinnedAddress(addresses: readonly ResolvedAddress[]): ResolvedAddress | null {
+  return addresses[0] ?? null;
+}
+
+function hostHeaderValue(parsed: URL): string {
+  return parsed.port.length === 0 ? parsed.host : `${parsed.hostname}:${parsed.port}`;
+}
+
+function headersToRecord(headers: HeadersInit | undefined, parsed: URL): Record<string, string> {
+  const normalized = new Headers(headers);
+  if (!normalized.has("host")) {
+    normalized.set("host", hostHeaderValue(parsed));
+  }
+  const result: Record<string, string> = {};
+  normalized.forEach((value, key) => {
+    result[key] = value;
+  });
+  return result;
+}
+
+async function requestBodyToUint8Array(
+  body: BodyInit | null | undefined,
+): Promise<
+  { ok: true; body: Uint8Array | null; contentType: string | null } | { ok: false; reason: string }
+> {
+  if (body === undefined || body === null) return { ok: true, body: null, contentType: null };
+  if (typeof body === "string") {
+    return { ok: true, body: new TextEncoder().encode(body), contentType: null };
+  }
+  if (body instanceof URLSearchParams) {
+    return {
+      ok: true,
+      body: new TextEncoder().encode(body.toString()),
+      contentType: "application/x-www-form-urlencoded;charset=UTF-8",
+    };
+  }
+  if (body instanceof Uint8Array) return { ok: true, body, contentType: null };
+  if (body instanceof ArrayBuffer) {
+    return { ok: true, body: new Uint8Array(body), contentType: null };
+  }
+  if (body instanceof Blob) {
+    return { ok: true, body: new Uint8Array(await body.arrayBuffer()), contentType: body.type };
+  }
+  return {
+    ok: false,
+    reason: "Pinned upstream fetch only supports buffered request bodies",
+  };
+}
+
+function pinnedLookup(
+  pinnedAddress: ResolvedAddress,
+): (
+  hostname: string,
+  options: { all?: boolean },
+  callback: (error: Error | null, address: string | ResolvedAddress[], family?: number) => void,
+) => void {
+  return (_hostname, options, callback) => {
+    if (options.all === true) {
+      callback(null, [pinnedAddress]);
+      return;
+    }
+    callback(null, pinnedAddress.address, pinnedAddress.family);
+  };
+}
+
+function requestOptionsForPinnedFetch(
+  parsed: URL,
+  init: RequestInit,
+  pinnedAddress: ResolvedAddress,
+  contentType: string | null,
+): RequestOptions {
+  const headers = headersToRecord(init.headers, parsed);
+  if (contentType !== null && !new Headers(headers).has("content-type")) {
+    headers["content-type"] = contentType;
+  }
+
+  return {
+    protocol: parsed.protocol,
+    hostname: parsed.hostname,
+    port: parsed.port,
+    path: `${parsed.pathname}${parsed.search}`,
+    method: init.method ?? "GET",
+    headers,
+    lookup: pinnedLookup(pinnedAddress),
+  };
+}
+
+function streamFromIncomingMessage(
+  incoming: IncomingMessage,
+  maxBytes: number | undefined,
+): ReadableStream<Uint8Array> {
+  let seenBytes = 0;
+  return new ReadableStream<Uint8Array>({
+    start(controller) {
+      incoming.on("data", (chunk: Uint8Array) => {
+        seenBytes += chunk.byteLength;
+        if (maxBytes !== undefined && seenBytes > maxBytes) {
+          incoming.destroy(
+            new Error(
+              `Upstream response exceeded configured limit of ${formatByteLimit(maxBytes)}`,
+            ),
+          );
+          controller.error(
+            new Error(
+              `Upstream response exceeded configured limit of ${formatByteLimit(maxBytes)}`,
+            ),
+          );
+          return;
+        }
+        controller.enqueue(chunk);
+      });
+      incoming.on("end", () => {
+        controller.close();
+      });
+      incoming.on("error", (error) => {
+        controller.error(error);
+      });
+    },
+    cancel() {
+      incoming.destroy();
+    },
+  });
+}
+
+async function cancelBody(response: Response): Promise<void> {
+  if (response.body === null) return;
+  await response.body.cancel().catch(() => undefined);
+}
+
+function responseWithByteLimit(response: Response, maxBytes: number): Response {
+  if (response.body === null) return response;
+
+  let seenBytes = 0;
+  const reader = response.body.getReader();
+  const body = new ReadableStream<Uint8Array>({
+    async pull(controller) {
+      const chunk = await reader.read();
+      if (chunk.done) {
+        controller.close();
+        return;
+      }
+
+      seenBytes += chunk.value.byteLength;
+      if (seenBytes > maxBytes) {
+        void reader.cancel().catch(() => undefined);
+        controller.error(
+          new Error(`Upstream response exceeded configured limit of ${formatByteLimit(maxBytes)}`),
+        );
+        return;
+      }
+
+      controller.enqueue(chunk.value);
+    },
+    cancel(reason) {
+      return reader.cancel(reason);
+    },
+  });
+
+  return new Response(body, {
+    status: response.status,
+    statusText: response.statusText,
+    headers: response.headers,
+  });
+}
+
+function shouldUseBudgetedSignal(init: RequestInit, deadlineMs: number | undefined): boolean {
+  return deadlineMs !== undefined || init.signal !== undefined;
+}
+
+async function pinnedNodeFetch(
+  rawUrl: string,
+  init: RequestInit,
+  options: SafeFetchOptions,
+  decision: Extract<ResolvedUpstreamUrlDecision, { allow: true }>,
+): Promise<SafeFetchResult> {
+  const parsed = new URL(rawUrl);
+  const pinnedAddress = choosePinnedAddress(decision.addresses);
+  if (pinnedAddress === null) {
+    return {
+      ok: false,
+      kind: "network-policy",
+      reason: `Could not pin a resolved upstream address for "${decision.hostname}"`,
+    };
+  }
+
+  const encodedBody = await requestBodyToUint8Array(init.body);
+  if (!encodedBody.ok) {
+    return { ok: false, kind: "resource-limit", reason: encodedBody.reason };
+  }
+
+  return await new Promise<SafeFetchResult>((resolve, reject) => {
+    const linkedSignal = init.signal ?? undefined;
+    const transport = parsed.protocol === "https:" ? requestHttps : requestHttp;
+    const requestOptions = requestOptionsForPinnedFetch(
+      parsed,
+      init,
+      pinnedAddress,
+      encodedBody.contentType,
+    );
+    const request = transport(requestOptions, (incoming) => {
+      const status = incoming.statusCode ?? 0;
+      if (isRedirectStatus(status)) {
+        incoming.resume();
+        const location = incoming.headers.location;
+        resolve({
+          ok: false,
+          kind: "redirect",
+          reason:
+            location === undefined
+              ? `Upstream returned blocked redirect HTTP ${String(status)}`
+              : `Upstream redirect to "${location}" is blocked; configure the final API URL directly`,
+        });
+        return;
+      }
+
+      const maxBytes = options.responseBytes;
+      const headers = new Headers();
+      for (const [key, value] of Object.entries(incoming.headers)) {
+        if (Array.isArray(value)) {
+          headers.set(key, value.join(", "));
+          continue;
+        }
+        if (value !== undefined) headers.set(key, value);
+      }
+
+      if (maxBytes !== undefined) {
+        const declaredBytes = parseContentLength(headers.get("content-length"));
+        if (declaredBytes !== null && declaredBytes > maxBytes) {
+          incoming.destroy();
+          resolve({
+            ok: false,
+            kind: "resource-limit",
+            reason: `Upstream declared ${String(declaredBytes)} response bytes, exceeding configured limit of ${formatByteLimit(maxBytes)}`,
+          });
+          return;
+        }
+      }
+
+      resolve({
+        ok: true,
+        response: new Response(streamFromIncomingMessage(incoming, maxBytes), {
+          status,
+          statusText: incoming.statusMessage,
+          headers,
+        }),
+      });
+    });
+
+    let timedOut = false;
+    let timeout: ReturnType<typeof setTimeout> | undefined;
+    const finishAbort = (): void => {
+      request.destroy(new DOMException("Safe fetch aborted", "AbortError"));
+    };
+
+    if (linkedSignal !== undefined) {
+      if (linkedSignal.aborted) {
+        finishAbort();
+      } else {
+        linkedSignal.addEventListener("abort", finishAbort, { once: true });
+      }
+    }
+
+    if (options.deadlineMs !== undefined) {
+      timeout = setTimeout(() => {
+        timedOut = true;
+        request.destroy(new DOMException("Safe fetch deadline exceeded", "TimeoutError"));
+      }, options.deadlineMs);
+    }
+
+    request.on("error", (error) => {
+      if (timeout !== undefined) clearTimeout(timeout);
+      if (linkedSignal !== undefined) {
+        linkedSignal.removeEventListener("abort", finishAbort);
+      }
+      if (timedOut) {
+        resolve({
+          ok: false,
+          kind: "timeout",
+          reason: `Upstream request exceeded configured deadline of ${String(options.deadlineMs)} ms`,
+        });
+        return;
+      }
+      reject(error);
+    });
+
+    request.on("close", () => {
+      if (timeout !== undefined) clearTimeout(timeout);
+      if (linkedSignal !== undefined) {
+        linkedSignal.removeEventListener("abort", finishAbort);
+      }
+    });
+
+    if (encodedBody.body !== null) request.write(encodedBody.body);
+    request.end();
+  });
+}
+
 export async function safeFetch(
   rawUrl: string,
   init: RequestInit = {},
   options: SafeFetchOptions = {},
 ): Promise<SafeFetchResult> {
-  const decision = await evaluateResolvedUpstreamUrl(rawUrl, options);
-  if (!decision.allow) {
-    return { ok: false, kind: "network-policy", reason: decision.reason };
+  const resolvedDecision = await evaluateResolvedUpstreamUrlForFetch(rawUrl, options);
+  if (!resolvedDecision.allow) {
+    return { ok: false, kind: "network-policy", reason: resolvedDecision.reason };
   }
 
-  const response = await fetch(rawUrl, { ...init, redirect: "manual" });
-  if (!isRedirectStatus(response.status)) return { ok: true, response };
-
-  if (response.body !== null) {
-    await response.body.cancel().catch(() => undefined);
+  if (isPinnedDnsEnabled(options)) {
+    return await pinnedNodeFetch(rawUrl, init, options, resolvedDecision);
   }
-  const location = response.headers.get("location");
-  return {
-    ok: false,
-    kind: "redirect",
-    reason:
-      location === null
-        ? `Upstream returned blocked redirect HTTP ${String(response.status)}`
-        : `Upstream redirect to "${location}" is blocked; configure the final API URL directly`,
+
+  const deadlineMs = options.deadlineMs;
+  const useBudgetedSignal = shouldUseBudgetedSignal(init, deadlineMs);
+  const controller = useBudgetedSignal ? new AbortController() : undefined;
+  let timedOut = false;
+  let timeout: ReturnType<typeof setTimeout> | undefined;
+  const linkedSignal = init.signal ?? undefined;
+  const forwardAbort = (): void => {
+    controller?.abort(linkedSignal?.reason);
   };
+
+  if (controller !== undefined && linkedSignal !== undefined) {
+    if (linkedSignal.aborted) {
+      forwardAbort();
+    } else {
+      linkedSignal.addEventListener("abort", forwardAbort, { once: true });
+    }
+  }
+
+  if (controller !== undefined && deadlineMs !== undefined) {
+    timeout = setTimeout(() => {
+      timedOut = true;
+      controller.abort(new DOMException("Safe fetch deadline exceeded", "TimeoutError"));
+    }, deadlineMs);
+  }
+
+  try {
+    const response = await fetch(rawUrl, {
+      ...init,
+      redirect: "manual",
+      signal: controller?.signal ?? init.signal,
+    });
+    if (!isRedirectStatus(response.status)) {
+      const maxBytes = options.responseBytes;
+      if (maxBytes === undefined) return { ok: true, response };
+
+      const declaredBytes = parseContentLength(response.headers.get("content-length"));
+      if (declaredBytes !== null && declaredBytes > maxBytes) {
+        await cancelBody(response);
+        return {
+          ok: false,
+          kind: "resource-limit",
+          reason: `Upstream declared ${String(declaredBytes)} response bytes, exceeding configured limit of ${formatByteLimit(maxBytes)}`,
+        };
+      }
+
+      return { ok: true, response: responseWithByteLimit(response, maxBytes) };
+    }
+
+    await cancelBody(response);
+    const location = response.headers.get("location");
+    return {
+      ok: false,
+      kind: "redirect",
+      reason:
+        location === null
+          ? `Upstream returned blocked redirect HTTP ${String(response.status)}`
+          : `Upstream redirect to "${location}" is blocked; configure the final API URL directly`,
+    };
+  } catch (error) {
+    if (timedOut) {
+      return {
+        ok: false,
+        kind: "timeout",
+        reason:
+          deadlineMs === undefined
+            ? "Upstream request exceeded its configured deadline"
+            : `Upstream request exceeded configured deadline of ${String(deadlineMs)} ms`,
+      };
+    }
+    return Promise.reject(error);
+  } finally {
+    if (timeout !== undefined) clearTimeout(timeout);
+    if (controller !== undefined && linkedSignal !== undefined) {
+      linkedSignal.removeEventListener("abort", forwardAbort);
+    }
+  }
 }