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,47 @@ 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));
+  return Math.max(256, Math.ceil((data.length * 3) / 4 / 512));
+}
+
+/**
+ * Char-based token estimate for an Anthropic-shaped request body. Base64 attachment
+ * payloads (image/document sources, wherever they nest — including 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.
+ */
+export function estimateClaudeRequestTokens(
+  raw: { system?: unknown; messages?: unknown; tools?: unknown },
+  modelId: string | undefined,
+): number {
+  let attachmentTokens = 0;
+  const stripAttachments = (value: unknown): string =>
+    JSON.stringify(value, (_key, entry: unknown) => {
+      if (entry && typeof entry === "object") {
+        const source = entry as { type?: unknown; data?: unknown };
+        if (source.type === "base64" && typeof source.data === "string") {
+          attachmentTokens += estimateBase64AttachmentTokens(source.data);
+          return { ...(entry as Record<string, unknown>), data: "" };
+        }
+      }
+      return entry;
+    });
+  const parts: string[] = [];
+  if (raw.system !== undefined) parts.push(typeof raw.system === "string" ? raw.system : stripAttachments(raw.system));
+  if (raw.messages !== undefined) parts.push(stripAttachments(raw.messages));
+  if (raw.tools !== undefined) parts.push(stripAttachments(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 +937,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,98 @@ 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 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);