From c8d7c09ca45a38827e22508c9afa5f14669f08a8 Mon Sep 17 00:00:00 2001
From: Krushna Korade <krishnacodes15@gmail.com>
Date: Wed, 14 Jan 2026 19:02:36 +0530
Subject: [PATCH] fix: PDF token estimation (#16494) (#16527)

Co-authored-by: Jack Wotherspoon <jackwoth@google.com>
---
 .../core/src/utils/tokenCalculation.test.ts   | 36 +++++++++++++++++++
 packages/core/src/utils/tokenCalculation.ts   | 11 ++++--
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/packages/core/src/utils/tokenCalculation.test.ts b/packages/core/src/utils/tokenCalculation.test.ts
index c6e54bc887..126ef7bac2 100644
--- a/packages/core/src/utils/tokenCalculation.test.ts
+++ b/packages/core/src/utils/tokenCalculation.test.ts
@@ -145,4 +145,40 @@ describe('calculateRequestTokenCount', () => {
 
     expect(count).toBe(3000);
   });
+
+  it('should use countTokens API for PDF requests', async () => {
+    vi.mocked(mockContentGenerator.countTokens).mockResolvedValue({
+      totalTokens: 5160,
+    });
+    const request = [
+      { inlineData: { mimeType: 'application/pdf', data: 'pdf_data' } },
+    ];
+
+    const count = await calculateRequestTokenCount(
+      request,
+      mockContentGenerator,
+      model,
+    );
+
+    expect(count).toBe(5160);
+    expect(mockContentGenerator.countTokens).toHaveBeenCalled();
+  });
+
+  it('should use fixed estimate for PDFs in fallback', async () => {
+    vi.mocked(mockContentGenerator.countTokens).mockRejectedValue(
+      new Error('API error'),
+    );
+    const request = [
+      { inlineData: { mimeType: 'application/pdf', data: 'large_pdf_data' } },
+    ];
+
+    const count = await calculateRequestTokenCount(
+      request,
+      mockContentGenerator,
+      model,
+    );
+
+    // PDF estimate: 25800 tokens (~100 pages at 258 tokens/page)
+    expect(count).toBe(25800);
+  });
 });
diff --git a/packages/core/src/utils/tokenCalculation.ts b/packages/core/src/utils/tokenCalculation.ts
index 06292bb925..cc2b83beb3 100644
--- a/packages/core/src/utils/tokenCalculation.ts
+++ b/packages/core/src/utils/tokenCalculation.ts
@@ -16,6 +16,9 @@ const ASCII_TOKENS_PER_CHAR = 0.25;
 const NON_ASCII_TOKENS_PER_CHAR = 1.3;
 // Fixed token estimate for images
 const IMAGE_TOKEN_ESTIMATE = 3000;
+// Fixed token estimate for PDFs (~100 pages at 258 tokens/page)
+// See: https://ai.google.dev/gemini-api/docs/document-processing
+const PDF_TOKEN_ESTIMATE = 25800;
 
 /**
  * Estimates token count for parts synchronously using a heuristic.
@@ -34,15 +37,19 @@ export function estimateTokenCountSync(parts: Part[]): number {
         }
       }
     } else {
-      // For images, we use a fixed safe estimate (3,000 tokens) covering
-      // up to 4K resolution on Gemini 3.
+      // For images and PDFs, we use fixed safe estimates:
+      // - Images: 3,000 tokens (covers up to 4K resolution on Gemini 3)
+      // - PDFs: 25,800 tokens (~100 pages at 258 tokens/page)
       // See: https://ai.google.dev/gemini-api/docs/vision#token_counting
+      // See: https://ai.google.dev/gemini-api/docs/document-processing
       const inlineData = 'inlineData' in part ? part.inlineData : undefined;
       const fileData = 'fileData' in part ? part.fileData : undefined;
       const mimeType = inlineData?.mimeType || fileData?.mimeType;
 
       if (mimeType?.startsWith('image/')) {
         totalTokens += IMAGE_TOKEN_ESTIMATE;
+      } else if (mimeType?.startsWith('application/pdf')) {
+        totalTokens += PDF_TOKEN_ESTIMATE;
       } else {
         // For other non-text parts (functionCall, functionResponse, etc.),
         // we fallback to the JSON string length heuristic.