feat(tracing): strip inline media from messages (#18413)

This is the functional portion addressing JS-1002.

Prior to truncating text messages for their byte length, any inline
base64-encoded media properties are filtered out. This allows the
message to possibly be included in the span, indicating to the user that
a media object was present, without overflowing the allotted buffer for
sending data.

If a media message is not removed, the fallback is still to simply
remove it if its overhead grows too large.

If any truncation occurs, then a `client.recordDroppedEvent()` message
is fired, with the reason `before_send`, category `attachment`, and the
count of messages affected.

Re JS-1002
Re GH-17810

Author: isaacs

packages/core/src/tracing/ai/messageTruncation.ts

@@ -1,3 +1,5 @@
+import { getClient } from '../../currentScopes';
+
 /**
  * Default maximum size in bytes for GenAI messages.
  * Messages exceeding this limit will be truncated.
@@ -12,20 +14,64 @@ type ContentMessage = {
   content: string;
 };
 
+/**
+ * Message format used by OpenAI and Anthropic APIs for media.
+ */
+type ContentArrayMessage = {
+  [key: string]: unknown;
+  content: {
+    [key: string]: unknown;
+    type: string;
+  }[];
+};
+
+/**
+ * Inline media content source, with a potentially very large base64
+ * blob or data: uri.
+ */
+type ContentMedia = Record<string, unknown> &
+  (
+    | {
+        media_type: string;
+        data: string;
+      }
+    | {
+        image_url: `data:${string}`;
+      }
+    | {
+        type: 'blob' | 'base64';
+        content: string;
+      }
+    | {
+        b64_json: string;
+      }
+    | {
+        uri: `data:${string}`;
+      }
+  );
+
 /**
  * Message format used by Google GenAI API.
  * Parts can be strings or objects with a text property.
  */
 type PartsMessage = {
   [key: string]: unknown;
-  parts: Array<string | { text: string }>;
+  parts: Array<TextPart | MediaPart>;
 };
 
 /**
  * A part in a Google GenAI message that contains text.
  */
 type TextPart = string | { text: string };
 
+/**
+ * A part in a Google GenAI that contains media.
+ */
+type MediaPart = {
+  type: string;
+  content: string;
+};
+
 /**
  * Calculate the UTF-8 byte length of a string.
  */
@@ -46,9 +92,10 @@ const jsonBytes = (value: unknown): number => {
  *
  * @param text - The string to truncate
  * @param maxBytes - Maximum byte length (UTF-8 encoded)
+ * @param onTruncate - Function called to record when truncation occurs
  * @returns Truncated string that fits within maxBytes
  */
-function truncateTextByBytes(text: string, maxBytes: number): string {
+function truncateTextByBytes(text: string, maxBytes: number, onTruncate: (count: number) => void): string {
   if (utf8Bytes(text) <= maxBytes) {
     return text;
   }
@@ -70,6 +117,7 @@ function truncateTextByBytes(text: string, maxBytes: number): string {
     }
   }
 
+  onTruncate(1);
   return bestFit;
 }
 
@@ -79,11 +127,12 @@ function truncateTextByBytes(text: string, maxBytes: number): string {
  *
  * @returns The text content
  */
-function getPartText(part: TextPart): string {
+function getPartText(part: TextPart | MediaPart): string {
   if (typeof part === 'string') {
     return part;
   }
-  return part.text;
+  if ('text' in part) return part.text;
+  return '';
 }
 
 /**
@@ -93,7 +142,7 @@ function getPartText(part: TextPart): string {
  * @param text - New text content
  * @returns New part with updated text
  */
-function withPartText(part: TextPart, text: string): TextPart {
+function withPartText(part: TextPart | MediaPart, text: string): TextPart {
   if (typeof part === 'string') {
     return text;
   }
@@ -112,6 +161,33 @@ function isContentMessage(message: unknown): message is ContentMessage {
   );
 }
 
+/**
+ * Check if a message has the OpenAI/Anthropic content array format.
+ */
+function isContentArrayMessage(message: unknown): message is ContentArrayMessage {
+  return message !== null && typeof message === 'object' && 'content' in message && Array.isArray(message.content);
+}
+
+/**
+ * Check if a content part is an OpenAI/Anthropic media source
+ */
+function isContentMedia(part: unknown): part is ContentMedia {
+  if (!part || typeof part !== 'object') return false;
+
+  return (
+    isContentMediaSource(part) ||
+    ('media_type' in part && typeof part.media_type === 'string' && 'data' in part) ||
+    ('image_url' in part && typeof part.image_url === 'string' && part.image_url.startsWith('data:')) ||
+    ('type' in part && (part.type === 'blob' || part.type === 'base64')) ||
+    'b64_json' in part ||
+    ('type' in part && 'result' in part && part.type === 'image_generation') ||
+    ('uri' in part && typeof part.uri === 'string' && part.uri.startsWith('data:'))
+  );
+}
+function isContentMediaSource(part: NonNullable<unknown>): boolean {
+  return 'type' in part && typeof part.type === 'string' && 'source' in part && isContentMedia(part.source);
+}
+
 /**
  * Check if a message has the Google GenAI parts format.
  */
@@ -132,17 +208,22 @@ function isPartsMessage(message: unknown): message is PartsMessage {
  * @param maxBytes - Maximum byte limit
  * @returns Array with truncated message, or empty array if it doesn't fit
  */
-function truncateContentMessage(message: ContentMessage, maxBytes: number): unknown[] {
+function truncateContentMessage(
+  message: ContentMessage,
+  maxBytes: number,
+  onTruncate: (count: number) => void,
+): unknown[] {
   // Calculate overhead (message structure without content)
   const emptyMessage = { ...message, content: '' };
   const overhead = jsonBytes(emptyMessage);
   const availableForContent = maxBytes - overhead;
 
   if (availableForContent <= 0) {
+    onTruncate(1);
     return [];
   }
 
-  const truncatedContent = truncateTextByBytes(message.content, availableForContent);
+  const truncatedContent = truncateTextByBytes(message.content, availableForContent, onTruncate);
   return [{ ...message, content: truncatedContent }];
 }
 
@@ -154,7 +235,7 @@ function truncateContentMessage(message: ContentMessage, maxBytes: number): unkn
  * @param maxBytes - Maximum byte limit
  * @returns Array with truncated message, or empty array if it doesn't fit
  */
-function truncatePartsMessage(message: PartsMessage, maxBytes: number): unknown[] {
+function truncatePartsMessage(message: PartsMessage, maxBytes: number, onTruncate: (count: number) => void): unknown[] {
   const { parts } = message;
 
   // Calculate overhead by creating empty text parts
@@ -163,11 +244,12 @@ function truncatePartsMessage(message: PartsMessage, maxBytes: number): unknown[
   let remainingBytes = maxBytes - overhead;
 
   if (remainingBytes <= 0) {
+    onTruncate(1);
     return [];
   }
 
   // Include parts until we run out of space
-  const includedParts: TextPart[] = [];
+  const includedParts: (TextPart | MediaPart)[] = [];
 
   for (const part of parts) {
     const text = getPartText(part);
@@ -179,7 +261,7 @@ function truncatePartsMessage(message: PartsMessage, maxBytes: number): unknown[
       remainingBytes -= textSize;
     } else if (includedParts.length === 0) {
       // First part doesn't fit: truncate it
-      const truncated = truncateTextByBytes(text, remainingBytes);
+      const truncated = truncateTextByBytes(text, remainingBytes, onTruncate);
       if (truncated) {
         includedParts.push(withPartText(part, truncated));
       }
@@ -190,7 +272,18 @@ function truncatePartsMessage(message: PartsMessage, maxBytes: number): unknown[
     }
   }
 
-  return includedParts.length > 0 ? [{ ...message, parts: includedParts }] : [];
+  /* c8 ignore start
+   * for type safety only, algorithm guarantees SOME text included */
+  if (includedParts.length <= 0) {
+    onTruncate(1);
+    return [];
+  } else {
+    /* c8 ignore stop */
+    if (includedParts.length < parts.length) {
+      onTruncate(1);
+    }
+    return [{ ...message, parts: includedParts }];
+  }
 }
 
 /**
@@ -204,23 +297,84 @@ function truncatePartsMessage(message: PartsMessage, maxBytes: number): unknown[
  * @param maxBytes - Maximum byte limit for the message
  * @returns Array containing the truncated message, or empty array if truncation fails
  */
-function truncateSingleMessage(message: unknown, maxBytes: number): unknown[] {
+function truncateSingleMessage(message: unknown, maxBytes: number, onTruncate: (count: number) => void): unknown[] {
+  /* c8 ignore start - unreachable */
   if (!message || typeof message !== 'object') {
+    onTruncate(1);
     return [];
   }
+  /* c8 ignore start - unreachable */
 
   if (isContentMessage(message)) {
-    return truncateContentMessage(message, maxBytes);
+    return truncateContentMessage(message, maxBytes, onTruncate);
   }
 
   if (isPartsMessage(message)) {
-    return truncatePartsMessage(message, maxBytes);
+    return truncatePartsMessage(message, maxBytes, onTruncate);
   }
 
   // Unknown message format: cannot truncate safely
+  onTruncate(1);
   return [];
 }
 
+const REMOVED_STRING = '<removed>';
+
+const MEDIA_FIELDS = ['image_url', 'data', 'content', 'b64_json', 'result', 'uri'] as const;
+
+function stripInlineMediaFromSingleMessage(part: ContentMedia, onTruncate: (count: number) => void): ContentMedia {
+  const strip = { ...part };
+  if (isContentMedia(strip.source)) {
+    strip.source = stripInlineMediaFromSingleMessage(strip.source, onTruncate);
+  }
+  for (const field of MEDIA_FIELDS) {
+    if (strip[field]) strip[field] = REMOVED_STRING;
+  }
+  return strip;
+}
+
+/**
+ * Strip the inline media from message arrays.
+ *
+ * This returns a stripped message. We do NOT want to mutate the data in place,
+ * because of course we still want the actual API/client to handle the media.
+ */
+function stripInlineMediaFromMessages(messages: unknown[], onTruncate: (count: number) => void): unknown[] {
+  let stripCount = 0;
+  const ont: (count: number) => void = count => (stripCount += count);
+  const stripped = messages.map(message => {
+    if (!!message && typeof message === 'object') {
+      if (isContentArrayMessage(message)) {
+        // eslint-disable-next-line no-param-reassign
+        message = {
+          ...message,
+          content: stripInlineMediaFromMessages(message.content, ont),
+        };
+      } else if ('content' in message && isContentMedia(message.content)) {
+        // eslint-disable-next-line no-param-reassign
+        message = {
+          ...message,
+          content: stripInlineMediaFromSingleMessage(message.content, ont),
+        };
+      }
+      if (isPartsMessage(message)) {
+        // eslint-disable-next-line no-param-reassign
+        message = {
+          ...message,
+          parts: stripInlineMediaFromMessages(message.parts, ont),
+        };
+      }
+      if (isContentMedia(message)) {
+        // eslint-disable-next-line no-param-reassign
+        message = stripInlineMediaFromSingleMessage(message, ont);
+      }
+    }
+    return message;
+  });
+  if (stripCount) onTruncate(stripCount);
+  return stripped;
+}
+
 /**
  * Truncate an array of messages to fit within a byte limit.
  *
@@ -240,12 +394,21 @@ function truncateSingleMessage(message: unknown, maxBytes: number): unknown[] {
  * // Returns [msg3, msg4] if they fit, or [msg4] if only it fits, etc.
  * ```
  */
-export function truncateMessagesByBytes(messages: unknown[], maxBytes: number): unknown[] {
+function truncateMessagesByBytes(
+  messages: unknown[],
+  maxBytes: number,
+  onTruncate: (count: number) => void,
+): unknown[] {
   // Early return for empty or invalid input
   if (!Array.isArray(messages) || messages.length === 0) {
     return messages;
   }
 
+  // strip inline media first. This will often get us below the threshold,
+  // while preserving human-readable information about messages sent.
+  // eslint-disable-next-line no-param-reassign
+  messages = stripInlineMediaFromMessages(messages, onTruncate);
+
   // Fast path: if all messages fit, return as-is
   const totalBytes = jsonBytes(messages);
   if (totalBytes <= maxBytes) {
@@ -275,8 +438,14 @@ export function truncateMessagesByBytes(messages: unknown[], maxBytes: number):
 
   // If no complete messages fit, try truncating just the newest message
   if (startIndex === messages.length) {
+    // we're truncating down to one message, so all others dropped.
+    if (messages.length !== 1) onTruncate(messages.length - 1);
     const newestMessage = messages[messages.length - 1];
-    return truncateSingleMessage(newestMessage, maxBytes);
+    return truncateSingleMessage(newestMessage, maxBytes, onTruncate);
+  }
+
+  if (startIndex !== 0) {
+    onTruncate(startIndex);
   }
 
   // Return the suffix that fits
@@ -292,7 +461,24 @@ export function truncateMessagesByBytes(messages: unknown[], maxBytes: number):
  * @returns Truncated array of messages
  */
 export function truncateGenAiMessages(messages: unknown[]): unknown[] {
-  return truncateMessagesByBytes(messages, DEFAULT_GEN_AI_MESSAGES_BYTE_LIMIT);
+  const { record, send } = getOnTruncate();
+  const truncated = truncateMessagesByBytes(messages, DEFAULT_GEN_AI_MESSAGES_BYTE_LIMIT, record);
+  send();
+  return truncated;
+}
+
+function getOnTruncate(): { record: (count: number) => void; send: () => void } {
+  let truncations = 0;
+  return {
+    record: (count: number) => {
+      truncations += count;
+    },
+    send: () => {
+      if (!truncations) return;
+      const client = getClient();
+      client?.recordDroppedEvent('before_send', 'attachment', truncations);
+    },
+  };
 }
 
 /**
@@ -302,5 +488,8 @@ export function truncateGenAiMessages(messages: unknown[]): unknown[] {
  * @returns Truncated string
  */
 export function truncateGenAiStringInput(input: string): string {
-  return truncateTextByBytes(input, DEFAULT_GEN_AI_MESSAGES_BYTE_LIMIT);
+  const { record, send } = getOnTruncate();
+  const truncated = truncateTextByBytes(input, DEFAULT_GEN_AI_MESSAGES_BYTE_LIMIT, record);
+  send();
+  return truncated;
 }