From f11d79a931d587ccf9c8717870a20bda60458b11 Mon Sep 17 00:00:00 2001
From: Abhi <43648792+abhipatel12@users.noreply.github.com>
Date: Sat, 13 Sep 2025 21:03:46 -0400
Subject: [PATCH] fix(stats): Use request model as fallback for stats logging
 (#8423)

---
 packages/core/src/code_assist/converter.test.ts   | 11 +++++++++++
 packages/core/src/code_assist/converter.ts        |  2 ++
 packages/core/src/core/loggingContentGenerator.ts | 14 ++++++++++----
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/packages/core/src/code_assist/converter.test.ts b/packages/core/src/code_assist/converter.test.ts
index 83dd418a27..c36c8424af 100644
--- a/packages/core/src/code_assist/converter.test.ts
+++ b/packages/core/src/code_assist/converter.test.ts
@@ -299,6 +299,17 @@ describe('converter', () => {
         codeAssistRes.response.automaticFunctionCallingHistory,
       );
     });
+
+    it('should handle modelVersion', () => {
+      const codeAssistRes: CaGenerateContentResponse = {
+        response: {
+          candidates: [],
+          modelVersion: 'gemini-2.5-pro',
+        },
+      };
+      const genaiRes = fromGenerateContentResponse(codeAssistRes);
+      expect(genaiRes.modelVersion).toEqual('gemini-2.5-pro');
+    });
   });
 
   describe('toContents', () => {
diff --git a/packages/core/src/code_assist/converter.ts b/packages/core/src/code_assist/converter.ts
index 147dcb01ef..78e743136f 100644
--- a/packages/core/src/code_assist/converter.ts
+++ b/packages/core/src/code_assist/converter.ts
@@ -80,6 +80,7 @@ interface VertexGenerateContentResponse {
   automaticFunctionCallingHistory?: Content[];
   promptFeedback?: GenerateContentResponsePromptFeedback;
   usageMetadata?: GenerateContentResponseUsageMetadata;
+  modelVersion?: string;
 }
 
 export interface CaCountTokenRequest {
@@ -137,6 +138,7 @@ export function fromGenerateContentResponse(
   out.automaticFunctionCallingHistory = inres.automaticFunctionCallingHistory;
   out.promptFeedback = inres.promptFeedback;
   out.usageMetadata = inres.usageMetadata;
+  out.modelVersion = inres.modelVersion;
   return out;
 }
 
diff --git a/packages/core/src/core/loggingContentGenerator.ts b/packages/core/src/core/loggingContentGenerator.ts
index 93a2a22deb..4d78d12a4a 100644
--- a/packages/core/src/core/loggingContentGenerator.ts
+++ b/packages/core/src/core/loggingContentGenerator.ts
@@ -114,7 +114,7 @@ export class LoggingContentGenerator implements ContentGenerator {
       const durationMs = Date.now() - startTime;
       this._logApiResponse(
         durationMs,
-        response.modelVersion || '',
+        response.modelVersion || req.model,
         userPromptId,
         response.usageMetadata,
         JSON.stringify(response),
@@ -143,13 +143,19 @@ export class LoggingContentGenerator implements ContentGenerator {
       throw error;
     }
 
-    return this.loggingStreamWrapper(stream, startTime, userPromptId);
+    return this.loggingStreamWrapper(
+      stream,
+      startTime,
+      userPromptId,
+      req.model,
+    );
   }
 
   private async *loggingStreamWrapper(
     stream: AsyncGenerator<GenerateContentResponse>,
     startTime: number,
     userPromptId: string,
+    model: string,
   ): AsyncGenerator<GenerateContentResponse> {
     const responses: GenerateContentResponse[] = [];
 
@@ -166,7 +172,7 @@ export class LoggingContentGenerator implements ContentGenerator {
       const durationMs = Date.now() - startTime;
       this._logApiResponse(
         durationMs,
-        responses[0]?.modelVersion || '',
+        responses[0]?.modelVersion || model,
         userPromptId,
         lastUsageMetadata,
         JSON.stringify(responses),
@@ -176,7 +182,7 @@ export class LoggingContentGenerator implements ContentGenerator {
       this._logApiError(
         durationMs,
         error,
-        responses[0]?.modelVersion || '',
+        responses[0]?.modelVersion || model,
         userPromptId,
       );
       throw error;