pi-advisor-flow 0.8.3-dev.3 → 0.9.0

package/src/image-validation.ts

@@ -0,0 +1,288 @@
+import { inflateSync } from "node:zlib";
+
+const PNG_SIGNATURE = Buffer.from("89504e470d0a1a0a", "hex");
+
+const isPng = (bytes: Buffer): boolean => {
+  if (bytes.length < 57 || !bytes.subarray(0, 8).equals(PNG_SIGNATURE)) {
+    return false;
+  }
+  let offset = 8;
+  let hasHeader = false;
+  const compressed: Buffer[] = [];
+  while (offset + 12 <= bytes.length) {
+    const length = bytes.readUInt32BE(offset);
+    if (length > bytes.length - offset - 12) {
+      return false;
+    }
+    const type = bytes.toString("ascii", offset + 4, offset + 8);
+    const data = offset + 8;
+    offset += length + 12;
+    if (!hasHeader) {
+      if (
+        type !== "IHDR" ||
+        length !== 13 ||
+        bytes.readUInt32BE(data) === 0 ||
+        bytes.readUInt32BE(data + 4) === 0
+      ) {
+        return false;
+      }
+      hasHeader = true;
+    } else if (type === "IDAT") {
+      compressed.push(bytes.subarray(data, data + length));
+    } else if (type === "IEND") {
+      if (length !== 0 || offset !== bytes.length || !compressed.length) {
+        return false;
+      }
+      try {
+        return (
+          inflateSync(Buffer.concat(compressed), {
+            maxOutputLength: 32 * 1024 * 1024,
+          }).length > 0
+        );
+      } catch {
+        return false;
+      }
+    } else if (type === "IHDR") {
+      return false;
+    }
+  }
+  return false;
+};
+
+const isFrameMarker = (marker: number) =>
+  [
+    0xc0, 0xc1, 0xc2, 0xc3, 0xc5, 0xc6, 0xc7, 0xc9, 0xca, 0xcb, 0xcd, 0xce,
+    0xcf,
+  ].includes(marker);
+
+const nextScanMarker = (bytes: Buffer, from: number): number => {
+  let offset = from;
+  while (offset < bytes.length - 1) {
+    if (bytes[offset] !== 0xff) {
+      offset += 1;
+      continue;
+    }
+    let next = offset + 1;
+    while (bytes[next] === 0xff) {
+      next += 1;
+    }
+    if (bytes[next] !== 0x00 && (bytes[next] < 0xd0 || bytes[next] > 0xd7)) {
+      return offset;
+    }
+    offset = next + 1;
+  }
+  return bytes.length;
+};
+
+interface JpegSegment {
+  end: number;
+  frame: boolean;
+  marker: number;
+  size: number;
+}
+
+const readJpegSegment = (
+  bytes: Buffer,
+  at: number
+): JpegSegment | undefined => {
+  if (bytes[at] !== 0xff) {
+    return;
+  }
+  let offset = at;
+  while (bytes[offset] === 0xff) {
+    offset += 1;
+  }
+  const marker = bytes[offset];
+  offset += 1;
+  if (marker === 0xd9) {
+    return { end: offset, frame: false, marker, size: 0 };
+  }
+  if (offset + 2 > bytes.length || marker === 0x00) {
+    return;
+  }
+  const size = bytes.readUInt16BE(offset);
+  if (size < 2 || offset + size > bytes.length) {
+    return;
+  }
+  const frame =
+    isFrameMarker(marker) &&
+    size >= 11 &&
+    bytes.readUInt16BE(offset + 3) > 0 &&
+    bytes.readUInt16BE(offset + 5) > 0;
+  return { end: offset + size, frame, marker, size };
+};
+
+const hasJpegEnvelope = (bytes: Buffer) =>
+  bytes.length >= 24 &&
+  bytes[0] === 0xff &&
+  bytes[1] === 0xd8 &&
+  bytes.at(-2) === 0xff &&
+  bytes.at(-1) === 0xd9;
+
+const isJpeg = (bytes: Buffer): boolean => {
+  if (!hasJpegEnvelope(bytes)) {
+    return false;
+  }
+  let offset = 2;
+  let frame = false;
+  let scan = false;
+  let quantization = false;
+  let entropyTable = false;
+  while (offset < bytes.length) {
+    const segment = readJpegSegment(bytes, offset);
+    if (!segment) {
+      return false;
+    }
+    const { end, marker, size } = segment;
+    if (marker === 0xd9) {
+      return frame && scan && end === bytes.length;
+    }
+    frame ||= segment.frame;
+    quantization ||= marker === 0xdb;
+    entropyTable ||= marker === 0xc4 || marker === 0xcc;
+    offset = end;
+    if (marker === 0xda) {
+      if (!(frame && quantization && entropyTable && size >= 6)) {
+        return false;
+      }
+      const next = nextScanMarker(bytes, offset);
+      if (next <= offset) {
+        return false;
+      }
+      scan = true;
+      offset = next;
+    }
+  }
+  return false;
+};
+
+const skipGifBlocks = (bytes: Buffer, at: number): number | undefined => {
+  let offset = at;
+  while (offset < bytes.length) {
+    const length = bytes[offset];
+    offset += 1;
+    if (length === 0) {
+      return offset;
+    }
+    if (offset + length > bytes.length) {
+      return;
+    }
+    offset += length;
+  }
+};
+
+const gifColorTableEnd = (bytes: Buffer, at: number, packed: number) =>
+  at + ((packed & 0x80) === 0 ? 0 : 3 * 2 ** ((packed & 7) + 1));
+
+const gifFrameEnd = (bytes: Buffer, at: number): number | undefined => {
+  if (
+    at + 9 > bytes.length ||
+    bytes.readUInt16LE(at + 4) === 0 ||
+    bytes.readUInt16LE(at + 6) === 0
+  ) {
+    return;
+  }
+  const data = gifColorTableEnd(bytes, at + 9, bytes[at + 8]);
+  if (
+    data + 2 >= bytes.length ||
+    bytes[data] < 2 ||
+    bytes[data] > 12 ||
+    bytes[data + 1] === 0
+  ) {
+    return;
+  }
+  return skipGifBlocks(bytes, data + 1);
+};
+
+const isGif = (bytes: Buffer): boolean => {
+  if (
+    bytes.length < 20 ||
+    !["GIF87a", "GIF89a"].includes(bytes.toString("ascii", 0, 6)) ||
+    bytes.readUInt16LE(6) === 0 ||
+    bytes.readUInt16LE(8) === 0
+  ) {
+    return false;
+  }
+  let offset = gifColorTableEnd(bytes, 13, bytes[10]);
+  let frame = false;
+  while (offset < bytes.length) {
+    const block = bytes[offset];
+    offset += 1;
+    if (block === 0x3b) {
+      return frame && offset === bytes.length;
+    }
+    if (block === 0x21 && offset < bytes.length) {
+      offset = skipGifBlocks(bytes, offset + 1) ?? bytes.length;
+    } else if (block === 0x2c) {
+      const end = gifFrameEnd(bytes, offset);
+      if (end === undefined) {
+        return false;
+      }
+      frame = true;
+      offset = end;
+    } else {
+      return false;
+    }
+  }
+  return false;
+};
+
+const validWebpFrame = (
+  bytes: Buffer,
+  type: string,
+  data: number,
+  length: number
+) => {
+  if (type === "VP8 ") {
+    return (
+      length >= 10 &&
+      bytes.toString("hex", data + 3, data + 6) === "9d012a" &&
+      (bytes.readUInt16LE(data + 6) & 0x3fff) > 0 &&
+      (bytes.readUInt16LE(data + 8) & 0x3fff) > 0
+    );
+  }
+  if (type === "VP8L") {
+    return length >= 5 && bytes[data] === 0x2f;
+  }
+  return type === "ANMF" && length >= 16;
+};
+
+const isWebp = (bytes: Buffer): boolean => {
+  if (
+    bytes.length < 30 ||
+    bytes.toString("ascii", 0, 4) !== "RIFF" ||
+    bytes.toString("ascii", 8, 12) !== "WEBP" ||
+    bytes.readUInt32LE(4) !== bytes.length - 8
+  ) {
+    return false;
+  }
+  let offset = 12;
+  let frame = false;
+  while (offset + 8 <= bytes.length) {
+    const type = bytes.toString("ascii", offset, offset + 4);
+    const length = bytes.readUInt32LE(offset + 4);
+    const data = offset + 8;
+    const end = data + length + (length % 2);
+    if (end > bytes.length) {
+      return false;
+    }
+    frame ||= validWebpFrame(bytes, type, data, length);
+    offset = end;
+  }
+  return frame && offset === bytes.length;
+};
+
+export const detectImageFormat = (bytes: Buffer): string | undefined => {
+  if (isPng(bytes)) {
+    return "image/png";
+  }
+  if (isJpeg(bytes)) {
+    return "image/jpeg";
+  }
+  if (isGif(bytes)) {
+    return "image/gif";
+  }
+  if (isWebp(bytes)) {
+    return "image/webp";
+  }
+};

package/src/images.ts

@@ -0,0 +1,151 @@
+import { createHash } from "node:crypto";
+
+import type { ImageContent, TextContent } from "@earendil-works/pi-ai/compat";
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+
+import type { AdvisorToolPolicies } from "./config/types.ts";
+import { contentParts, isRecordOf, isString } from "./content-utils.ts";
+import type { RecordValue } from "./content-utils.ts";
+import { detectImageFormat } from "./image-validation.ts";
+
+export const ADVISOR_IMAGE_MAX_BYTES = 4 * 1024 * 1024;
+export const ADVISOR_IMAGES_TOTAL_MAX_BYTES = 8 * 1024 * 1024;
+export const ADVISOR_IMAGES_MAX_COUNT = 4;
+
+export const imageFromBytes = (
+  bytes: Buffer,
+  mimeType?: string
+): ImageContent | undefined => {
+  if (!bytes.length || bytes.length > ADVISOR_IMAGE_MAX_BYTES) {
+    return;
+  }
+  const actual = detectImageFormat(bytes);
+  if (!actual || (mimeType && mimeType !== actual)) {
+    return;
+  }
+  return { data: bytes.toString("base64"), mimeType: actual, type: "image" };
+};
+
+export const imageFromPart = (part: RecordValue): ImageContent | undefined => {
+  if (
+    part.type !== "image" ||
+    !isString(part.data) ||
+    !isString(part.mimeType) ||
+    part.data.length > Math.ceil((ADVISOR_IMAGE_MAX_BYTES * 4) / 3) + 4 ||
+    !/^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/u.test(
+      part.data
+    )
+  ) {
+    return;
+  }
+  const bytes = Buffer.from(part.data, "base64");
+  if (bytes.toString("base64") !== part.data) {
+    return;
+  }
+  return imageFromBytes(bytes, part.mimeType);
+};
+
+export const imageMarker = (
+  part: RecordValue,
+  nonce = ""
+): string | undefined => {
+  if (part.type !== "image") {
+    return;
+  }
+  const image = imageFromPart(part);
+  if (!image) {
+    return "[Image omitted: unsupported format, invalid data, or over 4 MiB; pixels not reviewed]";
+  }
+  const id = createHash("sha256")
+    .update(nonce)
+    .update(image.data)
+    .digest("hex")
+    .slice(0, 24);
+  return `[Image ref=img_${id}; pixels reviewed only if attached below]`;
+};
+
+export interface SelectedImage {
+  image: ImageContent;
+  marker: string;
+}
+
+export interface ConversationImageCensus {
+  imagePartsSeen: number;
+  selected: SelectedImage[];
+}
+
+const eligibleImagePart = (part: RecordValue) =>
+  part.type === "image" && isString(part.data) && isString(part.mimeType);
+
+interface ImagePartCensus {
+  parts: number;
+  selected: SelectedImage[];
+}
+
+const entryImageCensus = (
+  content: string | (ImageContent | TextContent)[],
+  conversation: string,
+  nonce: string,
+  seen: Set<string>
+): ImagePartCensus => {
+  let parts = 0;
+  const selected: SelectedImage[] = [];
+  for (const part of contentParts(content)) {
+    if (!isRecordOf(part)) {
+      continue;
+    }
+    if (eligibleImagePart(part)) {
+      parts += 1;
+    }
+    const marker = imageMarker(part, nonce);
+    if (!marker || seen.has(marker) || !conversation.includes(marker)) {
+      continue;
+    }
+    const image = imageFromPart(part);
+    if (image) {
+      seen.add(marker);
+      selected.push({ image, marker });
+    }
+  }
+  return { parts, selected };
+};
+
+export const selectedConversationImages = (
+  ctx: ExtensionContext,
+  conversation: string,
+  policies: AdvisorToolPolicies,
+  selectedEntryIds: ReadonlySet<string> | undefined,
+  nonce: string
+): ConversationImageCensus => {
+  if (!conversation) {
+    return { imagePartsSeen: 0, selected: [] };
+  }
+  const entries = selectedEntryIds
+    ? ctx.sessionManager.buildContextEntries()
+    : ctx.sessionManager.getBranch();
+  const images: SelectedImage[] = [];
+  const seen = new Set<string>();
+  let imagePartsSeen = 0;
+  for (const [index, entry] of entries.entries()) {
+    if (
+      entry.type !== "message" ||
+      (selectedEntryIds && !selectedEntryIds.has(entry.id ?? String(index)))
+    ) {
+      continue;
+    }
+    const { message } = entry;
+    if (message.role !== "user" && message.role !== "toolResult") {
+      continue;
+    }
+    if (
+      message.role === "toolResult" &&
+      (policies[message.toolName] ?? "full") !== "full"
+    ) {
+      continue;
+    }
+    const census = entryImageCensus(message.content, conversation, nonce, seen);
+    imagePartsSeen += census.parts;
+    images.push(...census.selected);
+  }
+  return { imagePartsSeen, selected: images };
+};

package/src/scout-context.ts

@@ -39,6 +39,7 @@ const resolveCaps = (
   maxManifestBytes: number;
 } => ({
   currentInvocationId: options.currentInvocationId,
+  imageNonce: options.imageNonce ?? "",
   maxConversationChars: options.maxConversationChars,
   maxGroupBytes: options.maxGroupBytes ?? SCOUT_GROUP_MAX_BYTES,
   maxGroups: options.maxGroups ?? SCOUT_MANIFEST_MAX_GROUPS,

package/src/scout-curation.ts

@@ -20,10 +20,12 @@ export const curateAdvisorConversation = async (
   enabled = advisorScoutEnabledRef,
   runScout: typeof runAdvisorScout = runAdvisorScout,
   currentInvocationId?: string,
-  maxChars?: number
+  maxChars?: number,
+  imageNonce = ""
 ): Promise<{
   conversation: string;
   scout?: Exclude<ScoutOutcome, { cancelled: true }>;
+  selectedEntryIds?: string[];
 }> => {
   if (!enabled) {
     return { conversation: legacyConversation };
@@ -33,6 +35,7 @@ export const curateAdvisorConversation = async (
   }
   const built = buildScoutManifest(ctx, {
     currentInvocationId,
+    imageNonce,
     maxConversationChars: maxChars,
     maxManifestBytes: SCOUT_MANIFEST_MAX_BYTES,
   });
@@ -59,9 +62,12 @@ export const curateAdvisorConversation = async (
       ? signal.reason
       : new Error("Advisor operation cancelled during Scout.");
   }
-  let conversation = legacyConversation;
   if (outcome.ok) {
-    conversation =
+    const selected = new Set(outcome.selection.selectedIds);
+    const selectedEntryIds = built.manifest.groups
+      .filter((group) => group.required || selected.has(group.id))
+      .flatMap((group) => group.entryIds ?? []);
+    const conversation =
       maxChars === undefined
         ? outcome.conversation
         : reconstructScoutConversation(
@@ -70,6 +76,7 @@ export const curateAdvisorConversation = async (
             outcome.selection.synthesis,
             maxChars
           );
+    return { conversation, scout: outcome, selectedEntryIds };
   }
-  return { conversation, scout: outcome };
+  return { conversation: legacyConversation, scout: outcome };
 };

package/src/scout-groups.ts

@@ -76,6 +76,7 @@ const createGroup = (
 ): ScoutContextGroup => ({
   bytes: byteLength(content),
   content,
+  entryIds,
   id: stableId(originalIndex, entryIds, kind, content),
   kind,
   label: labelFor(kind, content),
@@ -101,6 +102,7 @@ const pendingInvocationDisclosure = (
   toolResultMaxBytes: number,
   policies: AdvisorToolPolicies,
   redact: boolean,
+  imageNonce: string,
   disclosed: string
 ) => {
   if (!isRecord(entry.message)) {
@@ -123,12 +125,15 @@ const pendingInvocationDisclosure = (
     toolResultMaxLines,
     toolResultMaxBytes,
     policies,
-    redact
+    redact,
+    true,
+    imageNonce
   );
 };
 
 export interface DisclosureCaps {
   currentInvocationId?: string;
+  imageNonce: string;
   policies: AdvisorToolPolicies;
   redact: boolean;
   toolResultMaxBytes: number;
@@ -191,7 +196,9 @@ const collectResults = (
       caps.toolResultMaxLines,
       caps.toolResultMaxBytes,
       caps.policies,
-      caps.redact
+      caps.redact,
+      true,
+      caps.imageNonce
     );
     if (resultText) {
       resultParts.push(resultText);
@@ -235,6 +242,7 @@ const missingOutcome = (
     caps.toolResultMaxBytes,
     caps.policies,
     caps.redact,
+    caps.imageNonce,
     disclosed
   );
   return {
@@ -450,7 +458,9 @@ export const buildGroups = (
       caps.toolResultMaxLines,
       caps.toolResultMaxBytes,
       caps.policies,
-      caps.redact
+      caps.redact,
+      true,
+      caps.imageNonce
     );
     if (!disclosed) {
       continue;

package/src/scout-types.ts

@@ -19,6 +19,7 @@ export type ScoutGroupKind =
 export interface ScoutContextGroup {
   bytes: number;
   content: string;
+  entryIds?: string[];
   id: string;
   kind: ScoutGroupKind;
   label: string;
@@ -51,6 +52,7 @@ export interface InvalidProtocolResult {
 
 export interface BuildScoutManifestOptions {
   currentInvocationId?: string;
+  imageNonce?: string;
   /** @deprecated Use maxManifestBytes for the Scout transport budget. */
   maxBytes?: number;
   /** Maximum reconstructed Advisor conversation characters. */

package/src/tools/consult-context.ts

@@ -1,3 +1,6 @@
+import { randomUUID } from "node:crypto";
+
+import type { ImageContent } from "@earendil-works/pi-ai/compat";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 
 import { readTrackedFiles, readUntrackedFiles } from "../attachments.ts";
@@ -7,12 +10,19 @@ import {
   advisorGitContextRef,
   advisorRedactSecretsRef,
   advisorScoutEnabledRef,
+  advisorToolPoliciesRef,
   advisorTrackedFileContentRef,
   advisorUntrackedContentRef,
   contextMaxCharsRef,
 } from "../config/state.ts";
 import { clampGitContextLevel, collectGitContext } from "../git.ts";
 import type { GitContextLevel } from "../git.ts";
+import { readImageFiles } from "../image-attachments.ts";
+import {
+  ADVISOR_IMAGES_MAX_COUNT,
+  ADVISOR_IMAGES_TOTAL_MAX_BYTES,
+  selectedConversationImages,
+} from "../images.ts";
 import { readProjectPreferences } from "../preferences.ts";
 import { redactAndCapText, redactSecrets } from "../redaction.ts";
 import { curateAdvisorConversation } from "../scout-curation.ts";
@@ -38,6 +48,7 @@ export interface ConsultationContextOptions {
   includeUntracked?: string[];
   onScout?: (event: ScoutLifecycleEvent) => void;
   signal?: AbortSignal;
+  supportsImages?: boolean;
 }
 
 export interface ConsultationContext {
@@ -47,6 +58,10 @@ export interface ConsultationContext {
   conversation: string;
   /** Redacted draft text, if a draft was supplied. */
   draftText?: string;
+  imageOmissions: number;
+  imagePartsSeen: number;
+  images: { image: ImageContent; label: string }[];
+  supportsImages: boolean;
   /** Redacted project preferences, if present. */
   preferences?: { bytes: number; text: string };
   /** Non-cancelling Scout outcome, when Scout ran. */
@@ -92,9 +107,12 @@ export const assembleConsultationContext = async (
     0,
     contextMaxCharsRef - changeText.length
   );
+  const imageNonce = randomUUID();
   const legacyConversation = advisorRequestConversation(
     ctx,
-    conversationBudget
+    conversationBudget,
+    true,
+    imageNonce
   );
   const curated = await curateAdvisorConversation(
     ctx,
@@ -104,7 +122,8 @@ export const assembleConsultationContext = async (
     advisorScoutEnabledRef,
     runAdvisorScout,
     options.currentInvocationId,
-    conversationBudget
+    conversationBudget,
+    imageNonce
   );
   const preferences = await readProjectPreferences(
     ctx,
@@ -135,12 +154,69 @@ export const assembleConsultationContext = async (
         untracked.reduce((sum, item) => sum + item.bytes, 0)
     )
   );
+  const supportsImages = options.supportsImages ?? false;
+  const untrackedImages = await readImageFiles(
+    ctx.cwd,
+    options.includeUntracked ?? [],
+    supportsImages && advisorUntrackedContentRef,
+    "untracked",
+    ADVISOR_IMAGES_TOTAL_MAX_BYTES,
+    ADVISOR_IMAGES_MAX_COUNT
+  );
+  const trackedImages = await readImageFiles(
+    ctx.cwd,
+    options.includeTracked ?? [],
+    supportsImages && advisorTrackedFileContentRef,
+    "tracked",
+    ADVISOR_IMAGES_TOTAL_MAX_BYTES -
+      untrackedImages.images.reduce((sum, item) => sum + item.bytes, 0),
+    ADVISOR_IMAGES_MAX_COUNT - untrackedImages.images.length
+  );
+  const images = [...untrackedImages.images, ...trackedImages.images].map(
+    (item) => ({
+      image: item.image,
+      label: `File ${JSON.stringify(item.path)}`,
+    })
+  );
+  let remainingBytes =
+    ADVISOR_IMAGES_TOTAL_MAX_BYTES -
+    [...untrackedImages.images, ...trackedImages.images].reduce(
+      (sum, item) => sum + item.bytes,
+      0
+    );
+  let imageOmissions = untrackedImages.omitted + trackedImages.omitted;
+  const census = selectedConversationImages(
+    ctx,
+    curated.conversation,
+    advisorToolPoliciesRef,
+    curated.scout?.ok === true
+      ? new Set<string>(curated.selectedEntryIds)
+      : undefined,
+    imageNonce
+  );
+  for (const item of census.selected) {
+    const bytes = Buffer.from(item.image.data, "base64").length;
+    if (
+      !supportsImages ||
+      images.length >= ADVISOR_IMAGES_MAX_COUNT ||
+      bytes > remainingBytes
+    ) {
+      imageOmissions += 1;
+      continue;
+    }
+    remainingBytes -= bytes;
+    images.push({ image: item.image, label: item.marker });
+  }
   return {
     changeText,
     conversation: curated.conversation,
     draftText,
+    imageOmissions,
+    imagePartsSeen: census.imagePartsSeen,
+    images,
     preferences,
     scout: curated.scout,
+    supportsImages,
     tracked,
     untracked,
   };