fix(claude): stop counting base64 attachments as raw characters in token estimates

src/server/claude-messages.ts

@@ -7,7 +7,7 @@
  * unchanged. The Responses output (SSE or JSON) is converted back to Anthropic shape.
  */
 import { FORWARD_HEADERS } from "../adapters/openai-responses";
-import { enforceAnthropicImageLimits } from "../adapters/anthropic-image-guard";
+import { enforceAnthropicImageLimits, sniffImageDimensions } from "../adapters/anthropic-image-guard";
 import { normalizeAnthropicImages } from "../adapters/anthropic-image-normalize";
 import { AnthropicRequestError, anthropicToResponsesTranslation, extractOcxEffortDirective, extractOcxRouteDirective, resolveInboundModel, type ClaudeCacheKeySource } from "../claude/inbound";
 import { resolveDesktop3pAlias } from "../claude/desktop-3p";
@@ -645,12 +645,7 @@ async function handleClaudeMessagesWithBudget(
     // accurate-usage adapters — the request-log merge is max(reported, estimate) and
     // would overwrite real usage (audit 133 R1#7).
     if (route.provider.adapter === "cursor" || route.provider.adapter === "kiro") {
-      const raw = anthropicBody as Rec;
-      const parts: string[] = [];
-      if (raw.system !== undefined) parts.push(typeof raw.system === "string" ? raw.system : JSON.stringify(raw.system));
-      if (raw.messages !== undefined) parts.push(JSON.stringify(raw.messages));
-      if (raw.tools !== undefined) parts.push(JSON.stringify(raw.tools));
-      logCtx.usageLogInputTokens = Math.max(1, estimateTokens(parts.join("\n"), requestedModel));
+      logCtx.usageLogInputTokens = estimateClaudeRequestTokens(anthropicBody as Rec, requestedModel);
     }
     // Effort safety valve (devlog 136 B6, audit 139 R2#2): opus-shaped aliases make
     // every routed model look like a reasoning model to Claude clients, so a forced
@@ -865,6 +860,68 @@ async function handleClaudeMessagesWithBudget(
 }
 
 /** Documented approximation: serialize system+messages+tools, run the char estimator. */
+/** Per-attachment token estimate for a base64 payload: real image dimensions when the
+ * header is sniffable (Anthropic prices images at ~pixels/750), else decoded bytes/512,
+ * min 256 — the same shape as the Kiro usage estimator (estimateKiroImageTokens). */
+function estimateBase64AttachmentTokens(data: string): number {
+  const dims = sniffImageDimensions(data);
+  if (dims) return Math.max(256, Math.ceil((dims.width * dims.height) / 750));
+  const unpadded = data.endsWith("==") ? data.length - 2 : data.endsWith("=") ? data.length - 1 : data.length;
+  return Math.max(256, Math.ceil(Math.floor((unpadded * 3) / 4) / 512));
+}
+
+/**
+ * Char-based token estimate for an Anthropic-shaped request body. Base64 attachment
+ * payloads (image/document blocks in message content, including blocks nested in
+ * tool_result.content) are counted as a bounded per-attachment estimate instead of raw
+ * characters: one 2MB screenshot is ~2.7M base64 chars, which the plain chars/token
+ * divide reports as hundreds of thousands of tokens versus a real cost around 1.6k.
+ * That breaks the >2x drift bound the estimator is held to (devlog 260711_claude_inbound
+ * 040 §3). Text and url sources are left in place and counted as characters, as is
+ * anything outside protocol content positions (tool_use.input, tool schemas).
+ */
+export function estimateClaudeRequestTokens(
+  raw: { system?: unknown; messages?: unknown; tools?: unknown },
+  modelId: string | undefined,
+): number {
+  let attachmentTokens = 0;
+  // Blank base64 payloads ONLY in protocol content positions: message content blocks and
+  // blocks nested in tool_result.content. tool_use.input and tool schemas can legitimately
+  // contain attachment-shaped JSON, and those bytes ARE serialized into function_call
+  // arguments / tool definitions for routed providers, so they must keep counting as text.
+  // system is text-only per the Anthropic protocol (no attachment sources), so it is
+  // stringified as-is.
+  const sanitizeBlock = (block: unknown): unknown => {
+    if (!block || typeof block !== "object") return block;
+    const b = block as Record<string, unknown>;
+    if (b.type === "image" || b.type === "document") {
+      const source = b.source as { type?: unknown; data?: unknown } | undefined;
+      if (source && typeof source === "object" && source.type === "base64" && typeof source.data === "string") {
+        attachmentTokens += estimateBase64AttachmentTokens(source.data);
+        return { ...b, source: { ...(source as Record<string, unknown>), data: "" } };
+      }
+      return block;
+    }
+    if (b.type === "tool_result" && Array.isArray(b.content)) {
+      return { ...b, content: (b.content as unknown[]).map(sanitizeBlock) };
+    }
+    return block;
+  };
+  const sanitizedMessages = (messages: unknown): unknown =>
+    Array.isArray(messages)
+      ? messages.map(message => {
+          if (!message || typeof message !== "object") return message;
+          const m = message as Record<string, unknown>;
+          return Array.isArray(m.content) ? { ...m, content: (m.content as unknown[]).map(sanitizeBlock) } : message;
+        })
+      : messages;
+  const parts: string[] = [];
+  if (raw.system !== undefined) parts.push(typeof raw.system === "string" ? raw.system : JSON.stringify(raw.system));
+  if (raw.messages !== undefined) parts.push(JSON.stringify(sanitizedMessages(raw.messages)));
+  if (raw.tools !== undefined) parts.push(JSON.stringify(raw.tools));
+  return Math.max(1, estimateTokens(parts.join("\n"), modelId) + attachmentTokens);
+}
+
 export async function handleClaudeCountTokens(req: Request, config: OcxConfig): Promise<Response> {
   const disabled = claudeInboundDisabled(config);
   if (disabled) return disabled;
@@ -901,11 +958,7 @@ export async function handleClaudeCountTokens(req: Request, config: OcxConfig):
   if (wantsNativePassthrough(req, config, model)) {
     return await anthropicNativePassthrough(req, config, { model, provider: "anthropic-native", surface: "claude" }, undefined, raw, "/v1/messages/count_tokens");
   }
-  const parts: string[] = [];
-  if (raw.system !== undefined) parts.push(typeof raw.system === "string" ? raw.system : JSON.stringify(raw.system));
-  if (raw.messages !== undefined) parts.push(JSON.stringify(raw.messages));
-  if (raw.tools !== undefined) parts.push(JSON.stringify(raw.tools));
-  const inputTokens = Math.max(1, estimateTokens(parts.join("\n"), model));
+  const inputTokens = estimateClaudeRequestTokens(raw, model);
   return new Response(JSON.stringify({ input_tokens: inputTokens }), {
     status: 200,
     headers: { "Content-Type": "application/json" },

tests/claude-messages-endpoint.test.ts

@@ -10,12 +10,14 @@ import { clearableDeadline } from "../src/lib/abort";
 import type { RequestLogContext } from "../src/server/request-log";
 import { startServer } from "../src/server";
 import {
+  estimateClaudeRequestTokens,
   fetchWithHeaderDeadline,
   handleClaudeMessages,
   readBoundedPassthroughBody,
   resolvePassthroughBodyGuard,
   tapAnthropicSseForLog,
 } from "../src/server/claude-messages";
+import { estimateTokens } from "../src/lib/token-estimate";
 import type { OcxConfig } from "../src/types";
 import { installIsolatedCodexHome, type IsolatedCodexHome } from "./helpers/isolated-codex-home";
 import { SERVER_BUDGET_MS } from "./helpers/test-budget";
@@ -935,6 +937,168 @@ test("count_tokens returns a positive estimate in the exact contract shape", asy
   }
 });
 
+/** Minimal PNG header (signature + IHDR) so the attachment sniffer can read real dimensions. */
+function countTokensPngBase64(width: number, height: number): string {
+  const u32be = (n: number): number[] => [(n >>> 24) & 0xff, (n >>> 16) & 0xff, (n >>> 8) & 0xff, n & 0xff];
+  const bytes = [
+    0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a,
+    ...u32be(13), 0x49, 0x48, 0x44, 0x52, // len + "IHDR"
+    ...u32be(width), ...u32be(height),
+    8, 6, 0, 0, 0, // bit depth, color type, etc.
+  ];
+  return Buffer.from(Uint8Array.from(bytes)).toString("base64");
+}
+
+test("count_tokens prices base64 attachments as attachments, not characters", async () => {
+  saveConfig(mockConfig("http://127.0.0.1:1/v1"));
+  const server = startServer(0);
+  try {
+    const data = "A".repeat(700_000); // ~512KB decoded; counting chars would report ~200k tokens
+    const response = await fetch(new URL("/v1/messages/count_tokens", server.url), {
+      method: "POST",
+      headers: { "content-type": "application/json" },
+      body: JSON.stringify({
+        model: "mock/test-model",
+        messages: [{
+          role: "user",
+          content: [
+            { type: "text", text: "what is in this screenshot?" },
+            { type: "image", source: { type: "base64", media_type: "image/png", data } },
+          ],
+        }],
+      }),
+    });
+    expect(response.status).toBe(200);
+    const json = await response.json() as { input_tokens: number };
+    // ceil(700000 * 3/4 / 512) = 1026 attachment tokens plus a small text remainder.
+    expect(json.input_tokens).toBeGreaterThanOrEqual(1026);
+    expect(json.input_tokens).toBeLessThan(2000);
+  } finally {
+    await server.stop(true);
+  }
+});
+
+test("estimateClaudeRequestTokens matches the plain char estimate for text-only bodies", () => {
+  const raw = {
+    system: "be brief",
+    messages: [{ role: "user", content: "count me please, this is a sentence" }],
+    tools: [{ name: "Read", input_schema: { type: "object" } }],
+  };
+  const parts = [raw.system, JSON.stringify(raw.messages), JSON.stringify(raw.tools)];
+  expect(estimateClaudeRequestTokens(raw, "m")).toBe(Math.max(1, estimateTokens(parts.join("\n"), "m")));
+});
+
+test("estimateClaudeRequestTokens prices sniffable images by pixel dimensions", () => {
+  const raw = {
+    messages: [{
+      role: "user",
+      content: [{ type: "image", source: { type: "base64", media_type: "image/png", data: countTokensPngBase64(1500, 2000) } }],
+    }],
+  };
+  const estimate = estimateClaudeRequestTokens(raw, "m");
+  // ceil(1500 * 2000 / 750) = 4000 attachment tokens plus the JSON skeleton.
+  expect(estimate).toBeGreaterThanOrEqual(4000);
+  expect(estimate).toBeLessThan(4100);
+});
+
+test("estimateClaudeRequestTokens strips base64 documents nested in tool_result content", () => {
+  const raw = {
+    messages: [{
+      role: "user",
+      content: [{
+        type: "tool_result",
+        tool_use_id: "t1",
+        content: [{ type: "document", source: { type: "base64", media_type: "application/pdf", data: "Q".repeat(400_000) } }],
+      }],
+    }],
+  };
+  const estimate = estimateClaudeRequestTokens(raw, "m");
+  // ceil(400000 * 3/4 / 512) = 586 tokens, nowhere near the ~114k a char count would report.
+  expect(estimate).toBeGreaterThanOrEqual(586);
+  expect(estimate).toBeLessThan(1000);
+});
+
+test("estimateClaudeRequestTokens does not charge base64 padding as payload bytes", () => {
+  // Exactly 131072 decoded bytes: 174764 base64 chars ending in "=". Counting the padding
+  // would yield 131073 bytes and charge 257 tokens instead of 256.
+  const data = Buffer.from(new Uint8Array(131_072)).toString("base64");
+  const raw = {
+    messages: [{ role: "user", content: [{ type: "image", source: { type: "base64", media_type: "image/png", data } }] }],
+  };
+  const stripped = {
+    messages: [{ role: "user", content: [{ type: "image", source: { type: "base64", media_type: "image/png", data: "" } }] }],
+  };
+
+  expect(estimateClaudeRequestTokens(raw, "m")).toBe(
+    Math.max(1, estimateTokens(JSON.stringify(stripped.messages), "m") + 256),
+  );
+});
+
+test("estimateClaudeRequestTokens keeps base64-shaped tool_use input counted as text", () => {
+  // tool_use.input is serialized into function_call arguments and sent upstream, so a
+  // {type:"base64", data} shape inside it is NOT an attachment and must count as text.
+  const raw = {
+    messages: [{
+      role: "assistant",
+      content: [{ type: "tool_use", id: "t1", name: "upload", input: { type: "base64", data: "B".repeat(40_000) } }],
+    }],
+  };
+
+  expect(estimateClaudeRequestTokens(raw, "m")).toBe(
+    Math.max(1, estimateTokens(JSON.stringify(raw.messages), "m")),
+  );
+});
+
+test("estimateClaudeRequestTokens leaves complete attachment-shaped tool_use input intact", () => {
+  // Even a full {type:"image", source:{type:"base64", data}} object inside tool_use.input
+  // is a tool argument, not an attachment: the translator replays it verbatim inside
+  // function_call arguments, so it must count at its serialized size.
+  const raw = {
+    messages: [{
+      role: "assistant",
+      content: [{
+        type: "tool_use",
+        id: "t1",
+        name: "upload_image",
+        input: { type: "image", source: { type: "base64", media_type: "image/png", data: "C".repeat(50_000) } },
+      }],
+    }],
+  };
+
+  expect(estimateClaudeRequestTokens(raw, "m")).toBe(
+    Math.max(1, estimateTokens(JSON.stringify(raw.messages), "m")),
+  );
+});
+
+test("estimateClaudeRequestTokens leaves attachment-shaped tool schemas intact", () => {
+  // Tool definitions are forwarded to routed providers; an attachment-shaped example in a
+  // schema is not an attachment either.
+  const raw = {
+    messages: [{ role: "user", content: "hi" }],
+    tools: [{
+      name: "upload",
+      input_schema: { type: "object" },
+      example: { type: "image", source: { type: "base64", media_type: "image/png", data: "D".repeat(30_000) } },
+    }],
+  };
+  const parts = [JSON.stringify(raw.messages), JSON.stringify(raw.tools)];
+
+  expect(estimateClaudeRequestTokens(raw, "m")).toBe(
+    Math.max(1, estimateTokens(parts.join("\n"), "m")),
+  );
+});
+
+test("estimateClaudeRequestTokens counts text-source documents as ordinary text", () => {
+  const text = "plain text document body ".repeat(40);
+  const raw = {
+    messages: [{
+      role: "user",
+      content: [{ type: "document", source: { type: "text", media_type: "text/plain", data: text } }],
+    }],
+  };
+  expect(estimateClaudeRequestTokens(raw, "m")).toBe(Math.max(1, estimateTokens(JSON.stringify(raw.messages), "m")));
+});
+
 test("claudeCode.enabled=false -> 403 permission_error on both routes", async () => {
   saveConfig(mockConfig("http://127.0.0.1:1/v1", { enabled: false }));
   const server = startServer(0);