From a7bfab4d96014a78da1f87f7cb5fb94b16fde1b8 Mon Sep 17 00:00:00 2001
From: Nick Salerni <nick.salerni@gmail.com>
Date: Fri, 5 Sep 2025 11:22:51 -0700
Subject: [PATCH] fix(token-usage): fix token usage recording for failed API
 calls (#7667)

Co-authored-by: Bryan Morgan <bryanmorgan@google.com>
---
 .../core/src/core/loggingContentGenerator.ts  | 21 ++++++++-----------
 1 file changed, 9 insertions(+), 12 deletions(-)

diff --git a/packages/core/src/core/loggingContentGenerator.ts b/packages/core/src/core/loggingContentGenerator.ts
index 916750b3cf..93a2a22deb 100644
--- a/packages/core/src/core/loggingContentGenerator.ts
+++ b/packages/core/src/core/loggingContentGenerator.ts
@@ -151,19 +151,26 @@ export class LoggingContentGenerator implements ContentGenerator {
     startTime: number,
     userPromptId: string,
   ): AsyncGenerator<GenerateContentResponse> {
-    let lastResponse: GenerateContentResponse | undefined;
     const responses: GenerateContentResponse[] = [];
 
     let lastUsageMetadata: GenerateContentResponseUsageMetadata | undefined;
     try {
       for await (const response of stream) {
         responses.push(response);
-        lastResponse = response;
         if (response.usageMetadata) {
           lastUsageMetadata = response.usageMetadata;
         }
         yield response;
       }
+      // Only log successful API response if no error occurred
+      const durationMs = Date.now() - startTime;
+      this._logApiResponse(
+        durationMs,
+        responses[0]?.modelVersion || '',
+        userPromptId,
+        lastUsageMetadata,
+        JSON.stringify(responses),
+      );
     } catch (error) {
       const durationMs = Date.now() - startTime;
       this._logApiError(
@@ -174,16 +181,6 @@ export class LoggingContentGenerator implements ContentGenerator {
       );
       throw error;
     }
-    const durationMs = Date.now() - startTime;
-    if (lastResponse) {
-      this._logApiResponse(
-        durationMs,
-        responses[0]?.modelVersion || '',
-        userPromptId,
-        lastUsageMetadata,
-        JSON.stringify(responses),
-      );
-    }
   }
 
   async countTokens(req: CountTokensParameters): Promise<CountTokensResponse> {