From c8ffda1616fdf46d50a0a7e1c32198e7ea2b8440 Mon Sep 17 00:00:00 2001
From: Vikhyath Mondreti <vikhyathvikku@gmail.com>
Date: Tue, 27 Jan 2026 19:16:54 -0800
Subject: [PATCH] fix(gemini): token count (#3039)

* fix(gemini): token count

* fix to include tool call tokens
---
 apps/sim/executor/handlers/agent/agent-handler.ts |  8 ++++++--
 apps/sim/executor/handlers/agent/types.ts         |  4 ++--
 apps/sim/providers/google/utils.ts                | 12 ++++++++----
 3 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/apps/sim/executor/handlers/agent/agent-handler.ts b/apps/sim/executor/handlers/agent/agent-handler.ts
index a22f7a8c4..007833d9c 100644
--- a/apps/sim/executor/handlers/agent/agent-handler.ts
+++ b/apps/sim/executor/handlers/agent/agent-handler.ts
@@ -936,8 +936,12 @@ export class AgentBlockHandler implements BlockHandler {
       systemPrompt: validMessages ? undefined : inputs.systemPrompt,
       context: validMessages ? undefined : stringifyJSON(messages),
       tools: formattedTools,
-      temperature: inputs.temperature != null ? Number(inputs.temperature) : undefined,
-      maxTokens: inputs.maxTokens != null ? Number(inputs.maxTokens) : undefined,
+      temperature:
+        inputs.temperature != null && inputs.temperature !== ''
+          ? Number(inputs.temperature)
+          : undefined,
+      maxTokens:
+        inputs.maxTokens != null && inputs.maxTokens !== '' ? Number(inputs.maxTokens) : undefined,
       apiKey: inputs.apiKey,
       azureEndpoint: inputs.azureEndpoint,
       azureApiVersion: inputs.azureApiVersion,
diff --git a/apps/sim/executor/handlers/agent/types.ts b/apps/sim/executor/handlers/agent/types.ts
index c3050f3a0..411b02a27 100644
--- a/apps/sim/executor/handlers/agent/types.ts
+++ b/apps/sim/executor/handlers/agent/types.ts
@@ -14,8 +14,8 @@ export interface AgentInputs {
   slidingWindowSize?: string // For message-based sliding window
   slidingWindowTokens?: string // For token-based sliding window
   // LLM parameters
-  temperature?: number
-  maxTokens?: number
+  temperature?: string
+  maxTokens?: string
   apiKey?: string
   azureEndpoint?: string
   azureApiVersion?: string
diff --git a/apps/sim/providers/google/utils.ts b/apps/sim/providers/google/utils.ts
index 724094784..c5040aab4 100644
--- a/apps/sim/providers/google/utils.ts
+++ b/apps/sim/providers/google/utils.ts
@@ -123,17 +123,21 @@ export function extractFunctionCallPart(candidate: Candidate | undefined): Part
 }
 
 /**
- * Converts usage metadata from SDK response to our format
+ * Converts usage metadata from SDK response to our format.
+ * Per Gemini docs, total = promptTokenCount + candidatesTokenCount + toolUsePromptTokenCount + thoughtsTokenCount
+ * We include toolUsePromptTokenCount in input and thoughtsTokenCount in output for correct billing.
  */
 export function convertUsageMetadata(
   usageMetadata: GenerateContentResponseUsageMetadata | undefined
 ): GeminiUsage {
-  const promptTokenCount = usageMetadata?.promptTokenCount ?? 0
-  const candidatesTokenCount = usageMetadata?.candidatesTokenCount ?? 0
+  const thoughtsTokenCount = usageMetadata?.thoughtsTokenCount ?? 0
+  const toolUsePromptTokenCount = usageMetadata?.toolUsePromptTokenCount ?? 0
+  const promptTokenCount = (usageMetadata?.promptTokenCount ?? 0) + toolUsePromptTokenCount
+  const candidatesTokenCount = (usageMetadata?.candidatesTokenCount ?? 0) + thoughtsTokenCount
   return {
     promptTokenCount,
     candidatesTokenCount,
-    totalTokenCount: usageMetadata?.totalTokenCount ?? promptTokenCount + candidatesTokenCount,
+    totalTokenCount: usageMetadata?.totalTokenCount ?? 0,
   }
 }