feat(billing): implement G1 AI credits overage flow with billing telemetry (#18590)

packages/core/src/code_assist/server.ts

@@ -21,7 +21,10 @@ import type {
   ConversationInteraction,
   StreamingLatency,
   RecordCodeAssistMetricsRequest,
+  GeminiUserTier,
+  Credits,
 } from './types.js';
+import { UserTierId } from './types.js';
 import type {
   ListExperimentsRequest,
   ListExperimentsResponse,
@@ -37,7 +40,15 @@ import type {
 import * as readline from 'node:readline';
 import { Readable } from 'node:stream';
 import type { ContentGenerator } from '../core/contentGenerator.js';
-import { UserTierId } from './types.js';
+import type { Config } from '../config/config.js';
+import {
+  G1_CREDIT_TYPE,
+  getG1CreditBalance,
+  isOverageEligibleModel,
+  shouldAutoUseCredits,
+} from '../billing/billing.js';
+import { logBillingEvent } from '../telemetry/loggers.js';
+import { CreditsUsedEvent } from '../telemetry/billingEvents.js';
 import type {
   CaCountTokenResponse,
   CaGenerateContentResponse,
@@ -72,6 +83,8 @@ export class CodeAssistServer implements ContentGenerator {
     readonly sessionId?: string,
     readonly userTier?: UserTierId,
     readonly userTierName?: string,
+    readonly paidTier?: GeminiUserTier,
+    readonly config?: Config,
   ) {}
 
   async generateContentStream(
@@ -80,6 +93,19 @@ export class CodeAssistServer implements ContentGenerator {
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
     role: LlmRole,
   ): Promise<AsyncGenerator<GenerateContentResponse>> {
+    const autoUse = this.config
+      ? shouldAutoUseCredits(
+          this.config.getBillingSettings().overageStrategy,
+          getG1CreditBalance(this.paidTier),
+        )
+      : false;
+    const modelIsEligible = isOverageEligibleModel(req.model);
+    const shouldEnableCredits = modelIsEligible && autoUse;
+
+    const enabledCreditTypes = shouldEnableCredits
+      ? ([G1_CREDIT_TYPE] as string[])
+      : undefined;
+
     const responses =
       await this.requestStreamingPost<CaGenerateContentResponse>(
         'streamGenerateContent',
@@ -88,6 +114,7 @@ export class CodeAssistServer implements ContentGenerator {
           userPromptId,
           this.projectId,
           this.sessionId,
+          enabledCreditTypes,
         ),
         req.config?.abortSignal,
       );
@@ -99,6 +126,9 @@ export class CodeAssistServer implements ContentGenerator {
     return (async function* (
       server: CodeAssistServer,
     ): AsyncGenerator<GenerateContentResponse> {
+      let totalConsumed = 0;
+      let lastRemaining = 0;
+
       for await (const response of responses) {
         if (isFirst) {
           streamingLatency.firstMessageLatency = formatProtoJsonDuration(
@@ -121,8 +151,38 @@ export class CodeAssistServer implements ContentGenerator {
           req.config?.abortSignal,
         );
 
+        if (response.consumedCredits) {
+          for (const credit of response.consumedCredits) {
+            if (credit.creditType === G1_CREDIT_TYPE && credit.creditAmount) {
+              totalConsumed += parseInt(credit.creditAmount, 10) || 0;
+            }
+          }
+        }
+        if (response.remainingCredits) {
+          // Sum all G1 credit entries for consistency with getG1CreditBalance
+          lastRemaining = response.remainingCredits.reduce((sum, credit) => {
+            if (credit.creditType === G1_CREDIT_TYPE && credit.creditAmount) {
+              return sum + (parseInt(credit.creditAmount, 10) || 0);
+            }
+            return sum;
+          }, 0);
+          server.updateCredits(response.remainingCredits);
+        }
+
         yield translatedResponse;
       }
+
+      // Emit credits used telemetry after the stream completes
+      if (totalConsumed > 0 && server.config) {
+        logBillingEvent(
+          server.config,
+          new CreditsUsedEvent(
+            req.model ?? 'unknown',
+            totalConsumed,
+            lastRemaining,
+          ),
+        );
+      }
     })(this);
   }
 
@@ -140,6 +200,7 @@ export class CodeAssistServer implements ContentGenerator {
         userPromptId,
         this.projectId,
         this.sessionId,
+        undefined,
       ),
       req.config?.abortSignal,
       GENERATE_CONTENT_RETRY_DELAY_IN_MILLISECONDS,
@@ -160,9 +221,29 @@ export class CodeAssistServer implements ContentGenerator {
       req.config?.abortSignal,
     );
 
+    if (response.remainingCredits) {
+      this.updateCredits(response.remainingCredits);
+    }
+
     return translatedResponse;
   }
 
+  private updateCredits(remainingCredits: Credits[]): void {
+    if (!this.paidTier) {
+      return;
+    }
+
+    // Replace the G1 credits entries with the latest remaining amounts.
+    // Non-G1 credits are preserved as-is.
+    const nonG1Credits = (this.paidTier.availableCredits ?? []).filter(
+      (c) => c.creditType !== G1_CREDIT_TYPE,
+    );
+    const updatedG1Credits = remainingCredits.filter(
+      (c) => c.creditType === G1_CREDIT_TYPE,
+    );
+    this.paidTier.availableCredits = [...nonG1Credits, ...updatedG1Credits];
+  }
+
   async onboardUser(
     req: OnboardUserRequest,
   ): Promise<LongRunningOperationResponse> {
@@ -192,6 +273,25 @@ export class CodeAssistServer implements ContentGenerator {
     }
   }
 
+  async refreshAvailableCredits(): Promise<void> {
+    if (!this.paidTier) {
+      return;
+    }
+    const res = await this.loadCodeAssist({
+      cloudaicompanionProject: this.projectId,
+      metadata: {
+        ideType: 'IDE_UNSPECIFIED',
+        platform: 'PLATFORM_UNSPECIFIED',
+        pluginType: 'GEMINI',
+        duetProject: this.projectId,
+      },
+      mode: 'HEALTH_CHECK',
+    });
+    if (res.paidTier?.availableCredits) {
+      this.paidTier.availableCredits = res.paidTier.availableCredits;
+    }
+  }
+
   async fetchAdminControls(
     req: FetchAdminControlsRequest,
   ): Promise<FetchAdminControlsResponse> {