fix(api): estimate inline base64 image tokens instead of counting the data URL as text so it does not falsely exceed the context window (#8368) (#8401)

changelog.d/fixes/8368-image-token-context.md

@@ -0,0 +1 @@
+- fix(api): estimate inline base64 image tokens instead of counting the data URL as text so it does not falsely exceed the context window (#8368)

open-sse/services/contextManager.ts

@@ -47,13 +47,134 @@ function getReserveTokensOverride(): number | null {
 // Rough chars-per-token ratio for quick estimation
 const CHARS_PER_TOKEN = 4;
 
+// Bounded per-image token budget used in place of measuring the raw base64
+// payload as text. In line with the owner's PoC (~1052 total for prompt +
+// 1 image) and litellm's calculate_img_tokens() default-count fast-path —
+// see #8368 research notes.
+const IMAGE_TOKEN_ESTIMATE = 1200;
+
+// Matches inline base64 data URLs, e.g. "data:image/png;base64,AAAA...".
+// Deliberately scoped to `data:image/...;base64,` so remote (http/https)
+// URLs and generic long base64 text strings stay on the text-estimation path.
+const INLINE_BASE64_IMAGE_RE = /^data:image\/[a-zA-Z0-9.+-]+;base64,/;
+
+function isInlineBase64ImageUrl(value: unknown): boolean {
+  return typeof value === "string" && INLINE_BASE64_IMAGE_RE.test(value);
+}
+
+// OpenAI chat.completions: { type: 'image_url', image_url: { url: 'data:...' } | 'data:...' }
+function matchesOpenAIImageUrlShape(node: Record<string, unknown>): boolean {
+  const imageUrl = node.image_url;
+  if (isInlineBase64ImageUrl(imageUrl)) return true;
+  return (
+    !!imageUrl &&
+    typeof imageUrl === "object" &&
+    isInlineBase64ImageUrl((imageUrl as Record<string, unknown>).url)
+  );
+}
+
+// AI SDK: { type: 'image', image: 'data:...' } (also covers Responses API's
+// { type: 'input_image', image_url: 'data:...' } via matchesOpenAIImageUrlShape above).
+function matchesAiSdkImageShape(node: Record<string, unknown>): boolean {
+  return node.type === "image" && isInlineBase64ImageUrl(node.image);
+}
+
+// Claude: { type: 'image', source: { type: 'base64', data: '...' } }
+function matchesClaudeSourceShape(node: Record<string, unknown>): boolean {
+  if (node.type !== "image") return false;
+  const source = node.source;
+  if (!source || typeof source !== "object") return false;
+  const src = source as Record<string, unknown>;
+  return src.type === "base64" && typeof src.data === "string";
+}
+
+// Gemini: { inlineData: { data: '...' } } | { inline_data: { data: '...' } }
+function matchesGeminiInlineDataShape(node: Record<string, unknown>): boolean {
+  const inlineData = node.inlineData ?? node.inline_data;
+  if (!inlineData || typeof inlineData !== "object") return false;
+  return typeof (inlineData as Record<string, unknown>).data === "string";
+}
+
 /**
- * Estimate token count from text length
+ * Detect the 5 documented inline-base64 image content-block shapes (see the
+ * shape-specific matchers above).
+ */
+function isInlineBase64ImageBlock(node: Record<string, unknown>): boolean {
+  return (
+    matchesOpenAIImageUrlShape(node) ||
+    matchesAiSdkImageShape(node) ||
+    matchesClaudeSourceShape(node) ||
+    matchesGeminiInlineDataShape(node)
+  );
+}
+
+/**
+ * Recursively walk a structured node, replacing every recognized inline
+ * base64 image block with a short placeholder (so its bulk is excluded from
+ * the char-count pass below) while accumulating a bounded per-image token
+ * cost. Returns the accumulated image token cost; the caller measures the
+ * placeholder-substituted structure with the normal char/4 heuristic.
+ *
+ * Non-image content (including remote image URLs and generic base64 text)
+ * is left untouched and continues to flow through the text-estimation path.
+ */
+function extractImageTokens(node: unknown, seen: Set<unknown>): { node: unknown; tokens: number } {
+  if (node === null || typeof node !== "object") {
+    return { node, tokens: 0 };
+  }
+  // Guard against cycles in structured request bodies.
+  if (seen.has(node)) return { node, tokens: 0 };
+  seen.add(node);
+
+  if (Array.isArray(node)) {
+    let tokens = 0;
+    const out = node.map((item) => {
+      const record =
+        item && typeof item === "object" && !Array.isArray(item)
+          ? (item as Record<string, unknown>)
+          : null;
+      if (record && isInlineBase64ImageBlock(record)) {
+        tokens += IMAGE_TOKEN_ESTIMATE;
+        return { __image_token_estimate__: IMAGE_TOKEN_ESTIMATE };
+      }
+      const result = extractImageTokens(item, seen);
+      tokens += result.tokens;
+      return result.node;
+    });
+    return { node: out, tokens };
+  }
+
+  const record = node as Record<string, unknown>;
+  if (isInlineBase64ImageBlock(record)) {
+    return { node: { __image_token_estimate__: IMAGE_TOKEN_ESTIMATE }, tokens: IMAGE_TOKEN_ESTIMATE };
+  }
+
+  let tokens = 0;
+  const out: Record<string, unknown> = {};
+  for (const [key, value] of Object.entries(record)) {
+    const result = extractImageTokens(value, seen);
+    out[key] = result.node;
+    tokens += result.tokens;
+  }
+  return { node: out, tokens };
+}
+
+/**
+ * Estimate token count from text length.
+ *
+ * Structured input is first walked for inline base64 image blocks (#8368):
+ * each recognized image block is substituted with a bounded per-image token
+ * budget instead of measuring its base64 payload as raw text, then the
+ * remainder of the structure is measured normally via the char/4 heuristic.
  */
 export function estimateTokens(text: string | object | null | undefined): number {
   if (!text) return 0;
-  const str = typeof text === "string" ? text : JSON.stringify(text);
-  return Math.ceil(str.length / CHARS_PER_TOKEN);
+  if (typeof text === "string") {
+    return Math.ceil(text.length / CHARS_PER_TOKEN);
+  }
+  const { node, tokens: imageTokens } = extractImageTokens(text, new Set());
+  const str = JSON.stringify(node);
+  return Math.ceil(str.length / CHARS_PER_TOKEN) + imageTokens;
 }
 
 /**

tests/unit/8368-image-token-context.test.ts

@@ -0,0 +1,120 @@
+import test from "node:test";
+import assert from "node:assert/strict";
+import { estimateTokens, getTokenLimit } from "../../open-sse/services/contextManager.ts";
+
+function makeFakePngBase64(approxBytes: number): string {
+  return Buffer.alloc(approxBytes, 65).toString("base64");
+}
+
+test("#8368: inline base64 PNG image_url is NOT counted as raw text (bounded image-token estimate)", () => {
+  const base64 = makeFakePngBase64(1_900_000); // ~1.9MB, matches issue repro
+  const messages = [
+    { role: "user", content: "Please describe this image." },
+    {
+      role: "user",
+      content: [
+        { type: "input_text", text: "Please describe this image." },
+        { type: "input_image", image_url: `data:image/png;base64,${base64}` },
+      ],
+    },
+  ];
+  const estimated = estimateTokens(messages);
+  assert.ok(
+    estimated < 5000,
+    `BUG #8368 reproduced: image-bearing message estimated at ${estimated} tokens (limit ${getTokenLimit(
+      "codex"
+    )})`
+  );
+});
+
+test("#8368: plain text estimation is unaffected by the image-token fix (control)", () => {
+  const text = "a".repeat(4000); // 4000 chars => ~1000 tokens at CHARS_PER_TOKEN=4
+  const estimated = estimateTokens(text);
+  assert.equal(estimated, 1000);
+});
+
+test("#8368: OpenAI chat.completions image_url object shape is bounded", () => {
+  const base64 = makeFakePngBase64(500_000);
+  const messages = [
+    {
+      role: "user",
+      content: [
+        { type: "text", text: "What is in this image?" },
+        { type: "image_url", image_url: { url: `data:image/jpeg;base64,${base64}` } },
+      ],
+    },
+  ];
+  const estimated = estimateTokens(messages);
+  assert.ok(estimated < 5000, `expected bounded estimate, got ${estimated}`);
+});
+
+test("#8368: Claude source.base64 image block shape is bounded", () => {
+  const base64 = makeFakePngBase64(500_000);
+  const messages = [
+    {
+      role: "user",
+      content: [
+        { type: "text", text: "Describe this." },
+        {
+          type: "image",
+          source: { type: "base64", media_type: "image/png", data: base64 },
+        },
+      ],
+    },
+  ];
+  const estimated = estimateTokens(messages);
+  assert.ok(estimated < 5000, `expected bounded estimate, got ${estimated}`);
+});
+
+test("#8368: Gemini inlineData image block shape is bounded", () => {
+  const base64 = makeFakePngBase64(500_000);
+  const messages = [
+    {
+      role: "user",
+      parts: [{ text: "Describe this." }, { inlineData: { mimeType: "image/png", data: base64 } }],
+    },
+  ];
+  const estimated = estimateTokens(messages);
+  assert.ok(estimated < 5000, `expected bounded estimate, got ${estimated}`);
+});
+
+test("#8368: multiple images accumulate a bounded sum, not one flat cap", () => {
+  const base64 = makeFakePngBase64(200_000);
+  const oneImageMessages = [
+    { role: "user", content: [{ type: "image_url", image_url: { url: `data:image/png;base64,${base64}` } }] },
+  ];
+  const threeImageMessages = [
+    {
+      role: "user",
+      content: [
+        { type: "image_url", image_url: { url: `data:image/png;base64,${base64}` } },
+        { type: "image_url", image_url: { url: `data:image/png;base64,${base64}` } },
+        { type: "image_url", image_url: { url: `data:image/png;base64,${base64}` } },
+      ],
+    },
+  ];
+  const one = estimateTokens(oneImageMessages);
+  const three = estimateTokens(threeImageMessages);
+  assert.ok(three > one, `expected 3 images to cost more than 1 (one=${one}, three=${three})`);
+  assert.ok(three < one * 4, `expected roughly linear scaling, got one=${one} three=${three}`);
+});
+
+test("#8368: remote http(s) image URLs are unaffected (still measured as text)", () => {
+  const messages = [
+    {
+      role: "user",
+      content: [{ type: "image_url", image_url: { url: "https://example.com/cat.png" } }],
+    },
+  ];
+  const estimated = estimateTokens(messages);
+  // Should just be the JSON-length/4 heuristic for the short URL string, not near-zero
+  // and not a bounded image-token substitute — remote URLs stay on the text path.
+  assert.ok(estimated > 0 && estimated < 100, `expected small text-based estimate, got ${estimated}`);
+});
+
+test("#8368: generic long base64 text (not an image field) still uses the text path", () => {
+  const genericBase64 = makeFakePngBase64(100_000);
+  const estimated = estimateTokens(genericBase64);
+  const expectedTextEstimate = Math.ceil(genericBase64.length / 4);
+  assert.equal(estimated, expectedTextEstimate);
+});