@fitchmultz/pi-mcp-adapter 5.0.0

package/mcp-output-guard.ts

@@ -0,0 +1,412 @@
+import { randomBytes } from "node:crypto";
+import { mkdir, mkdtemp, writeFile } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join, resolve } from "node:path";
+import type { ContentBlock, McpSettings } from "./types.ts";
+
+export const DEFAULT_MCP_OUTPUT_MAX_BYTES = 50 * 1024;
+export const DEFAULT_MCP_OUTPUT_MAX_LINES = 2000;
+export const DEFAULT_MCP_DETAILS_MAX_BYTES = 16 * 1024;
+
+const CONTENT_SUMMARY_LIMIT = 20;
+const KEY_PREVIEW_LIMIT = 20;
+const KEY_MAX_CHARS = 120;
+
+type Recordish = Record<string, unknown>;
+
+export interface McpOutputGuardDetails {
+  truncated: true;
+  originalBytes: number;
+  returnedBytes: number;
+  originalLines: number;
+  returnedLines: number;
+  /** Number of image content blocks returned untouched alongside the truncated text. */
+  imageBlocksPassedThrough?: number;
+  fullOutputPath?: string;
+  writeError?: string;
+}
+
+export interface McpResultSummary {
+  omitted: true;
+  reason: string;
+  isError: boolean;
+  contentBlocks: number;
+  contentSummary: Array<Record<string, unknown>>;
+  structuredContentSummary?: Record<string, unknown>;
+  meta?: Record<string, unknown>;
+  extraFields?: Array<Record<string, unknown>>;
+  rawResultBytes: number;
+  fullResultPath?: string;
+  resultWriteError?: string;
+}
+
+export interface McpOutputGuardOptions {
+  enabled?: boolean;
+  outputDirectory?: string;
+  prefix?: string;
+  suffix?: string;
+  emptyTextFallback?: string;
+  maxBytes?: number;
+  maxLines?: number;
+  detailsMaxBytes?: number;
+  /**
+   * Raw MCP result to expose as details.mcpResult. Kept raw when its JSON
+   * fits detailsMaxBytes (or when the guard is disabled); otherwise replaced
+   * with a compact summary and spilled to a temp file. Omit for call sites
+   * whose details never carried the raw result (e.g. resource reads).
+   */
+  rawMcpResult?: unknown;
+}
+
+export interface GuardedMcpOutput {
+  content: ContentBlock[];
+  outputGuard?: McpOutputGuardDetails;
+  mcpResult?: unknown;
+}
+
+export function resolveMcpOutputGuardOptions(settings?: McpSettings, outputDirectory?: string): Pick<McpOutputGuardOptions, "enabled" | "maxBytes" | "maxLines" | "detailsMaxBytes" | "outputDirectory"> {
+  const configured = settings?.outputGuard;
+  const tuning = typeof configured === "object" && configured !== null ? configured : undefined;
+  return {
+    enabled: envKillSwitch("MCP_OUTPUT_GUARD") ?? configured !== false,
+    ...(outputDirectory !== undefined ? { outputDirectory } : {}),
+    maxBytes: positiveInt(tuning?.maxBytes) ?? DEFAULT_MCP_OUTPUT_MAX_BYTES,
+    maxLines: positiveInt(tuning?.maxLines) ?? DEFAULT_MCP_OUTPUT_MAX_LINES,
+    detailsMaxBytes: positiveInt(tuning?.detailsMaxBytes) ?? DEFAULT_MCP_DETAILS_MAX_BYTES,
+  };
+}
+
+/** Spread helper for tool-result details: includes mcpResult/outputGuard only when present. */
+export function guardedMcpDetails(guarded: GuardedMcpOutput): Record<string, unknown> {
+  return {
+    ...(guarded.mcpResult !== undefined ? { mcpResult: guarded.mcpResult } : {}),
+    ...(guarded.outputGuard ? { outputGuard: guarded.outputGuard } : {}),
+  };
+}
+
+/**
+ * Bound model-facing MCP output. Text output is capped at maxBytes/maxLines and
+ * spilled to a temp file when oversized. Image blocks pass through untouched —
+ * they are delivered to the provider as native image content, not text context.
+ */
+export async function guardMcpOutput(
+  content: ContentBlock[],
+  options: McpOutputGuardOptions = {},
+): Promise<GuardedMcpOutput> {
+  const maxBytes = options.maxBytes ?? DEFAULT_MCP_OUTPUT_MAX_BYTES;
+  const maxLines = options.maxLines ?? DEFAULT_MCP_OUTPUT_MAX_LINES;
+  const detailsMaxBytes = options.detailsMaxBytes ?? DEFAULT_MCP_DETAILS_MAX_BYTES;
+  const prefix = options.prefix ?? "";
+  const suffix = options.suffix ?? "";
+
+  const normalizedContent = withEmptyTextFallback(
+    content.length > 0
+      ? sanitizeContent(content)
+      : [{ type: "text" as const, text: options.emptyTextFallback ?? "(empty result)" }],
+    options.emptyTextFallback,
+  );
+
+  if (options.enabled === false) {
+    return {
+      content: addAffixes(normalizedContent, prefix, suffix),
+      ...(options.rawMcpResult !== undefined ? { mcpResult: options.rawMcpResult } : {}),
+    };
+  }
+
+  const imageBlocks = normalizedContent.filter((block) => block.type === "image");
+  const textOutput = normalizedContent
+    .filter((block) => block.type === "text")
+    .map((block) => (block as { text: string }).text)
+    .join("\n");
+  const composedOutput = `${prefix}${textOutput}${suffix}`;
+  const stats = textStats(composedOutput);
+
+  let guardedContent: ContentBlock[] = addAffixes(normalizedContent, prefix, suffix);
+  let outputGuard: McpOutputGuardDetails | undefined;
+
+  if (stats.bytes > maxBytes || stats.lines > maxLines) {
+    const { path: fullOutputPath, error: writeError } = await saveArtifact("output", composedOutput, options.outputDirectory);
+    const notice = formatTruncationNotice(stats, fullOutputPath, writeError);
+    const previewBudget = reserveBudget(maxBytes, maxLines, notice);
+    const preview = truncateHead(composedOutput, previewBudget.maxBytes, previewBudget.maxLines);
+    const finalText = `${preview.content}\n\n${notice}`;
+    const finalStats = textStats(finalText);
+
+    guardedContent = [{ type: "text" as const, text: finalText }, ...imageBlocks];
+    outputGuard = {
+      truncated: true,
+      originalBytes: stats.bytes,
+      returnedBytes: finalStats.bytes,
+      originalLines: stats.lines,
+      returnedLines: finalStats.lines,
+      ...(imageBlocks.length > 0 ? { imageBlocksPassedThrough: imageBlocks.length } : {}),
+      ...(fullOutputPath !== undefined ? { fullOutputPath } : {}),
+      ...(writeError !== undefined ? { writeError } : {}),
+    };
+  }
+
+  const mcpResult = options.rawMcpResult === undefined
+    ? undefined
+    : await boundMcpResult(options.rawMcpResult, detailsMaxBytes, options.outputDirectory);
+
+  return {
+    content: guardedContent,
+    ...(outputGuard ? { outputGuard } : {}),
+    ...(mcpResult !== undefined ? { mcpResult } : {}),
+  };
+}
+
+function sanitizeContent(content: ContentBlock[]): ContentBlock[] {
+  return content.map((block) => {
+    if (block.type !== "image") return block;
+    const mimeType = typeof block.mimeType === "string" && block.mimeType.trim()
+      ? block.mimeType.trim().slice(0, 100)
+      : "image/png";
+    return { ...block, mimeType };
+  });
+}
+
+function withEmptyTextFallback(content: ContentBlock[], fallback: string | undefined): ContentBlock[] {
+  if (!fallback) return content;
+  const textOutput = content
+    .filter((block) => block.type === "text")
+    .map((block) => (block as { text: string }).text)
+    .join("\n");
+  if (textOutput) return content;
+  return [{ type: "text", text: fallback }, ...content.filter((block) => block.type === "image")];
+}
+
+function addAffixes(content: ContentBlock[], prefix: string, suffix: string): ContentBlock[] {
+  if (!prefix && !suffix) return content;
+  const next: ContentBlock[] = [...content];
+
+  if (prefix) {
+    const index = next.findIndex((block) => block.type === "text");
+    const block = next[index];
+    if (block?.type === "text") {
+      next[index] = { ...block, text: `${prefix}${block.text}` };
+    } else {
+      next.unshift({ type: "text", text: prefix });
+    }
+  }
+
+  if (suffix) {
+    let index = -1;
+    for (let i = next.length - 1; i >= 0; i--) {
+      if (next[i]?.type === "text") {
+        index = i;
+        break;
+      }
+    }
+    const block = next[index];
+    if (block?.type === "text") {
+      next[index] = { ...block, text: `${block.text}${suffix}` };
+    } else {
+      next.push({ type: "text", text: suffix });
+    }
+  }
+
+  return next;
+}
+
+function reserveBudget(maxBytes: number, maxLines: number, notice: string): { maxBytes: number; maxLines: number } {
+  const noticeStats = textStats(`\n\n${notice}`);
+  return {
+    maxBytes: Math.max(0, maxBytes - noticeStats.bytes),
+    maxLines: Math.max(0, maxLines - noticeStats.lines),
+  };
+}
+
+function truncateHead(text: string, maxBytes: number, maxLines: number): { content: string; bytes: number; lines: number } {
+  const lines = text.split("\n");
+  const output: string[] = [];
+  let bytes = 0;
+
+  for (const line of lines) {
+    if (output.length >= maxLines) break;
+    const separatorBytes = output.length > 0 ? 1 : 0;
+    const lineBytes = byteLength(line);
+    if (bytes + separatorBytes + lineBytes > maxBytes) {
+      const remaining = maxBytes - bytes - separatorBytes;
+      if (remaining > 0) {
+        output.push(truncateStringToBytes(line, remaining));
+      }
+      break;
+    }
+    output.push(line);
+    bytes += separatorBytes + lineBytes;
+  }
+
+  const content = output.join("\n");
+  const stats = textStats(content);
+  return { content, bytes: stats.bytes, lines: stats.lines };
+}
+
+function truncateStringToBytes(value: string, maxBytes: number): string {
+  if (byteLength(value) <= maxBytes) return value;
+  const buffer = Buffer.from(value, "utf8");
+  let end = Number.isFinite(maxBytes) ? Math.max(0, Math.floor(maxBytes)) : 0;
+  while (end > 0 && (buffer.readUInt8(end) & 0xc0) === 0x80) end--;
+  return buffer.subarray(0, end).toString("utf8");
+}
+
+function formatTruncationNotice(
+  stats: { bytes: number; lines: number },
+  fullOutputPath: string | undefined,
+  writeError: string | undefined,
+): string {
+  const base = `[MCP text output truncated: original ${stats.lines.toLocaleString()} lines / ${formatSize(stats.bytes)}.`;
+  if (fullOutputPath) {
+    return `${base} Full text saved to: ${fullOutputPath} — use read with offset/limit or grep to inspect.]`;
+  }
+  return `${base} Full output could not be saved: ${writeError ?? "unknown error"}]`;
+}
+
+/**
+ * Bound details.mcpResult: keep the raw result when its JSON fits within
+ * detailsMaxBytes; otherwise replace it with a compact summary and spill the
+ * raw JSON to a temp file.
+ */
+async function boundMcpResult(result: unknown, detailsMaxBytes: number, outputDirectory?: string): Promise<unknown> {
+  const raw = safeStringify(result);
+  const rawBytes = byteLength(raw);
+  if (rawBytes <= detailsMaxBytes) return result;
+  return summarizeMcpResult(result, raw, rawBytes, outputDirectory);
+}
+
+async function summarizeMcpResult(result: unknown, raw: string, rawBytes: number, outputDirectory?: string): Promise<McpResultSummary> {
+  const { path: fullResultPath, error: resultWriteError } = await saveArtifact("mcp-result", raw, outputDirectory);
+
+  const record = asRecord(result);
+  const content = Array.isArray(record?.content) ? record.content : [];
+  const summary: McpResultSummary = {
+    omitted: true,
+    reason: "Raw MCP result exceeded the details size limit and was replaced with this summary to keep session context bounded.",
+    isError: record?.isError === true,
+    contentBlocks: content.length,
+    contentSummary: summarizeContent(content),
+    rawResultBytes: rawBytes,
+    ...(fullResultPath !== undefined ? { fullResultPath } : {}),
+    ...(resultWriteError !== undefined ? { resultWriteError } : {}),
+  };
+
+  if (record && "structuredContent" in record) {
+    summary.structuredContentSummary = summarizeValue(record.structuredContent);
+  }
+  if (record && "_meta" in record) {
+    summary.meta = summarizeValue(record._meta);
+  }
+  if (record) {
+    const standard = new Set(["content", "isError", "structuredContent", "_meta"]);
+    const extraFields = Object.keys(record)
+      .filter((key) => !standard.has(key))
+      .slice(0, KEY_PREVIEW_LIMIT)
+      .map((key) => ({ key: truncateKey(key), type: typeof record[key], estimatedBytes: estimateValueBytes(record[key]), omitted: true }));
+    if (extraFields.length > 0) summary.extraFields = extraFields;
+  }
+
+  return summary;
+}
+
+function summarizeContent(content: unknown[]): Array<Record<string, unknown>> {
+  const summaries: Array<Record<string, unknown>> = content.slice(0, CONTENT_SUMMARY_LIMIT).map((block) => {
+    const record = asRecord(block);
+    if (!record) return { type: typeof block, omitted: true };
+    if (record.type === "text") {
+      const text = typeof record.text === "string" ? record.text : "";
+      return { type: "text", bytes: byteLength(text), lines: textStats(text).lines, textOmitted: true };
+    }
+    if (record.type === "image") {
+      const data = typeof record.data === "string" ? record.data : "";
+      return { type: "image", mimeType: typeof record.mimeType === "string" ? record.mimeType : undefined, dataBytes: byteLength(data), dataOmitted: true };
+    }
+    return { type: typeof record.type === "string" ? record.type : "unknown", estimatedBytes: estimateValueBytes(record), omitted: true };
+  });
+  if (content.length > CONTENT_SUMMARY_LIMIT) {
+    summaries.push({ type: "omitted", count: content.length - CONTENT_SUMMARY_LIMIT });
+  }
+  return summaries;
+}
+
+function summarizeValue(value: unknown): Record<string, unknown> {
+  const record = asRecord(value);
+  if (!record) {
+    return { type: value === null ? "null" : typeof value, estimatedBytes: estimateValueBytes(value), omitted: true };
+  }
+  const keys = Object.keys(record);
+  return {
+    type: Array.isArray(value) ? "array" : "object",
+    estimatedBytes: estimateValueBytes(value),
+    keyCount: keys.length,
+    keysPreview: keys.slice(0, KEY_PREVIEW_LIMIT).map(truncateKey),
+    omitted: true,
+  };
+}
+
+function estimateValueBytes(value: unknown, depth = 0): number {
+  if (value === null || value === undefined) return 0;
+  if (typeof value === "string") return byteLength(value);
+  if (typeof value === "number" || typeof value === "boolean" || typeof value === "bigint") return byteLength(String(value));
+  const record = asRecord(value);
+  if (!record || depth >= 2) return 0;
+  const values = Array.isArray(value) ? value.slice(0, KEY_PREVIEW_LIMIT) : Object.values(record).slice(0, KEY_PREVIEW_LIMIT);
+  return values.reduce((total, item) => total + estimateValueBytes(item, depth + 1), 0);
+}
+
+function truncateKey(key: string): string {
+  return key.length <= KEY_MAX_CHARS ? key : `${key.slice(0, KEY_MAX_CHARS - 1)}…`;
+}
+
+async function saveArtifact(kind: string, text: string, outputDirectory?: string): Promise<{ path?: string; error?: string }> {
+  try {
+    const parent = outputDirectory === undefined ? tmpdir() : resolve(outputDirectory);
+    if (outputDirectory !== undefined) await mkdir(parent, { recursive: true, mode: 0o700 });
+    const dir = await mkdtemp(join(parent, "pi-mcp-output-"));
+    const path = join(dir, `${kind}-${randomBytes(4).toString("hex")}.txt`);
+    await writeFile(path, text, { encoding: "utf8", mode: 0o600 });
+    return { path };
+  } catch (error) {
+    return { error: error instanceof Error ? error.message : String(error) };
+  }
+}
+
+function asRecord(value: unknown): Recordish | undefined {
+  return typeof value === "object" && value !== null ? value as Recordish : undefined;
+}
+
+function safeStringify(value: unknown): string {
+  try {
+    // The output guard measures and spills raw MCP results; it does not render this JSON for the model.
+    return JSON.stringify(value);
+  } catch {
+    return String(value);
+  }
+}
+
+function textStats(text: string): { bytes: number; lines: number } {
+  return { bytes: byteLength(text), lines: text.length === 0 ? 0 : text.split("\n").length };
+}
+
+function byteLength(text: string): number {
+  return Buffer.byteLength(text, "utf8");
+}
+
+function positiveInt(value: unknown): number | undefined {
+  if (typeof value !== "number" || !Number.isFinite(value)) return undefined;
+  const integer = Math.floor(value);
+  return integer > 0 ? integer : undefined;
+}
+
+function envKillSwitch(name: string): boolean | undefined {
+  const value = process.env[name]?.trim().toLowerCase();
+  if (!value) return undefined;
+  if (["0", "false", "no", "off"].includes(value)) return false;
+  if (["1", "true", "yes", "on"].includes(value)) return true;
+  return undefined;
+}
+
+function formatSize(bytes: number): string {
+  if (bytes < 1024) return `${bytes} B`;
+  if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)} KiB`;
+  return `${(bytes / (1024 * 1024)).toFixed(1)} MiB`;
+}